Pass JSON document class through settings

This updates the JSON codec settings to include the document class
to use during deserialization. This is primarily intended to allow
customization of the discriminator.

Author: JordonPhillips

packages/smithy-json/src/smithy_json/__init__.py

@@ -9,11 +9,13 @@
 from smithy_core.serializers import ShapeSerializer
 from smithy_core.types import TimestampFormat
 
-from ._private import JSONSettings as _JSONSettings
 from ._private.deserializers import JSONShapeDeserializer as _JSONShapeDeserializer
+from ._private.documents import JSONDocument
 from ._private.serializers import JSONShapeSerializer as _JSONShapeSerializer
+from .settings import JSONSettings
 
 __version__: str = importlib.metadata.version("smithy-json")
+__all__ = ("JSONCodec", "JSONDocument", "JSONSettings")
 
 
 class JSONCodec(Codec):
@@ -25,6 +27,7 @@ def __init__(
         use_timestamp_format: bool = True,
         default_timestamp_format: TimestampFormat = TimestampFormat.DATE_TIME,
         default_namespace: str | None = None,
+        document_class: type[JSONDocument] = JSONDocument,
     ) -> None:
         """Initializes a JSONCodec.
 
@@ -34,11 +37,16 @@ def __init__(
             `smithy.api#timestampFormat` trait, if present.
         :param default_timestamp_format: The default timestamp format to use if the
             `smithy.api#timestampFormat` trait is not enabled or not present.
+        :param default_namespace: The default namespace to use when determining a
+            document's discriminator.
+        :param document_class: The document class to deserialize to.
         """
-        self._settings = _JSONSettings(
+        self._settings = JSONSettings(
             use_json_name=use_json_name,
             use_timestamp_format=use_timestamp_format,
             default_timestamp_format=default_timestamp_format,
+            default_namespace=default_namespace,
+            document_class=document_class,
         )
 
     @property

packages/smithy-json/src/smithy_json/_private/__init__.py

@@ -1,21 +1,10 @@
 #  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 #  SPDX-License-Identifier: Apache-2.0
-
-from dataclasses import dataclass
 from typing import Protocol, runtime_checkable
 
-from smithy_core.types import TimestampFormat
-
 
 @runtime_checkable
 class Flushable(Protocol):
     """A protocol for objects that can be flushed."""
 
     def flush(self) -> None: ...
-
-
-@dataclass
-class JSONSettings:
-    use_json_name: bool = True
-    use_timestamp_format: bool = True
-    default_timestamp_format: TimestampFormat = TimestampFormat.DATE_TIME

packages/smithy-json/src/smithy_json/_private/deserializers.py

@@ -18,8 +18,7 @@
 from smithy_core.traits import JSONNameTrait, TimestampFormatTrait
 from smithy_core.types import TimestampFormat
 
-from . import JSONSettings
-from .documents import JSONDocument
+from ..settings import JSONSettings
 
 # TODO: put these type hints in a pyi somewhere. There here because ijson isn't
 # typed.
@@ -89,11 +88,9 @@ def peek(self) -> JSONParseEvent:
 
 
 class JSONShapeDeserializer(ShapeDeserializer):
-    def __init__(
-        self, source: BytesReader, *, settings: JSONSettings | None = None
-    ) -> None:
+    def __init__(self, source: BytesReader, settings: JSONSettings) -> None:
         self._stream = BufferedParser(ijson.parse(source))
-        self._settings = settings or JSONSettings()
+        self._settings = settings
 
         # A mapping of json name to member name for each shape. Since the deserializer
         # is shared and we don't know which shapes will be deserialized, this is
@@ -158,7 +155,9 @@ def read_string(self, schema: Schema) -> str:
     def read_document(self, schema: Schema) -> Document:
         start = next(self._stream)
         if start.type not in ("start_map", "start_array"):
-            return JSONDocument(start.value, schema=schema, settings=self._settings)
+            return self._settings.document_class(
+                value=start.value, schema=schema, settings=self._settings
+            )
 
         end_type = "end_map" if start.type == "start_map" else "end_array"
         builder = cast(TypedObjectBuilder, ObjectBuilder())
@@ -168,7 +167,9 @@ def read_document(self, schema: Schema) -> Document:
         ).path != start.path or event.type != end_type:
             builder.event(event.type, event.value)
 
-        return JSONDocument(builder.value, schema=schema, settings=self._settings)
+        return self._settings.document_class(
+            value=builder.value, schema=schema, settings=self._settings
+        )
 
     def read_timestamp(self, schema: Schema) -> datetime.datetime:
         format = self._settings.default_timestamp_format

packages/smithy-json/src/smithy_json/_private/documents.py

@@ -13,23 +13,22 @@
 from smithy_core.traits import JSONNameTrait, TimestampFormatTrait
 from smithy_core.utils import expect_type
 
-from . import JSONSettings
+from ..settings import JSONSettings
 
 
 class JSONDocument(Document):
-    _schema: Schema
-    _json_names: dict[str, str]
+    _discriminator: ShapeID | None = None
 
     def __init__(
         self,
         value: DocumentValue | dict[str, "Document"] | list["Document"],
+        settings: JSONSettings | None = None,
         *,
         schema: Schema = DOCUMENT,
-        settings: JSONSettings | None = None,
     ) -> None:
         super().__init__(value, schema=schema)
-        self._settings = settings or JSONSettings()
-        self._json_names = {}
+        self._settings = settings or JSONSettings(document_class=type(self))
+        self._json_names: dict[str, str] = {}
 
         if self._settings.use_json_name and schema.shape_type in (
             ShapeType.STRUCTURE,
@@ -39,12 +38,6 @@ def __init__(
                 if json_name := member_schema.get_trait(JSONNameTrait):
                     self._json_names[json_name.value] = member_name
 
-    @property
-    def discriminator(self) -> ShapeID:
-        if self._type is ShapeType.MAP:
-            return ShapeID(self.as_map()["__type"].as_string())
-        return super().discriminator
-
     def as_blob(self) -> bytes:
         return b64decode(expect_type(str, self._value))
 

packages/smithy-json/src/smithy_json/_private/serializers.py

@@ -21,18 +21,17 @@
 from smithy_core.traits import JSONNameTrait, TimestampFormatTrait
 from smithy_core.types import TimestampFormat
 
-from . import Flushable, JSONSettings
+from ..settings import JSONSettings
+from . import Flushable
 
 _INF: float = float("inf")
 _NEG_INF: float = float("-inf")
 
 
 class JSONShapeSerializer(ShapeSerializer):
-    def __init__(
-        self, sink: BytesWriter, *, settings: JSONSettings | None = None
-    ) -> None:
+    def __init__(self, sink: BytesWriter, settings: JSONSettings) -> None:
         self._stream = StreamingJSONEncoder(sink)
-        self._settings = settings or JSONSettings()
+        self._settings = settings
 
     def begin_struct(
         self, schema: "Schema"

packages/smithy-json/src/smithy_json/settings.py

@@ -0,0 +1,31 @@
+#  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#  SPDX-License-Identifier: Apache-2.0
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+from smithy_core.types import TimestampFormat
+
+if TYPE_CHECKING:
+    from ._private.documents import JSONDocument
+
+
+@dataclass(slots=True)
+class JSONSettings:
+    """Settings for the JSON codec."""
+
+    document_class: type["JSONDocument"]
+    """The document class to deserialize to."""
+
+    use_json_name: bool = True
+    """Whether the codec should use `smithy.api#jsonName` trait, if present."""
+
+    use_timestamp_format: bool = True
+    """Whether the codec should use the `smithy.api#timestampFormat` trait, if
+    present."""
+
+    default_timestamp_format: TimestampFormat = TimestampFormat.DATE_TIME
+    """The default timestamp format to use if the `smithy.api#timestampFormat` trait is
+    not enabled or not present."""
+
+    default_namespace: str | None = None
+    """The default namespace to use when determining a document's discriminator."""

packages/smithy-json/tests/unit/test_deserializers.py

@@ -9,12 +9,13 @@
     BIG_DECIMAL,
     BLOB,
     BOOLEAN,
+    DOCUMENT,
     FLOAT,
     INTEGER,
     STRING,
     TIMESTAMP,
 )
-from smithy_json import JSONCodec
+from smithy_json import JSONCodec, JSONDocument
 
 from . import (
     JSON_SERDE_CASES,
@@ -88,3 +89,13 @@ def _read_optional_map(k: str, d: ShapeDeserializer):
         assert actual_value == expected_value
     else:
         assert actual == expected
+
+
+class CustomDocument(JSONDocument):
+    pass
+
+
+def test_uses_custom_document() -> None:
+    codec = JSONCodec(document_class=CustomDocument)
+    actual = codec.create_deserializer(b'{"foo": "bar"}').read_document(DOCUMENT)
+    assert isinstance(actual, CustomDocument)