Use shared settings for JSON codec

Author: JordonPhillips

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

@@ -9,6 +9,7 @@
 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.serializers import JSONShapeSerializer as _JSONShapeSerializer
 
@@ -18,15 +19,12 @@
 class JSONCodec(Codec):
     """A codec for converting shapes to/from JSON."""
 
-    _use_json_name: bool
-    _use_timestamp_format: bool
-    _default_timestamp_format: TimestampFormat
-
     def __init__(
         self,
         use_json_name: bool = True,
         use_timestamp_format: bool = True,
         default_timestamp_format: TimestampFormat = TimestampFormat.DATE_TIME,
+        default_namespace: str | None = None,
     ) -> None:
         """Initializes a JSONCodec.
 
@@ -37,28 +35,20 @@ def __init__(
         :param default_timestamp_format: The default timestamp format to use if the
             `smithy.api#timestampFormat` trait is not enabled or not present.
         """
-        self._use_json_name = use_json_name
-        self._use_timestamp_format = use_timestamp_format
-        self._default_timestamp_format = default_timestamp_format
+        self._settings = _JSONSettings(
+            use_json_name=use_json_name,
+            use_timestamp_format=use_timestamp_format,
+            default_timestamp_format=default_timestamp_format,
+        )
 
     @property
     def media_type(self) -> str:
         return "application/json"
 
     def create_serializer(self, sink: BytesWriter) -> "ShapeSerializer":
-        return _JSONShapeSerializer(
-            sink,
-            use_json_name=self._use_json_name,
-            use_timestamp_format=self._use_timestamp_format,
-            default_timestamp_format=self._default_timestamp_format,
-        )
+        return _JSONShapeSerializer(sink, settings=self._settings)
 
     def create_deserializer(self, source: bytes | BytesReader) -> "ShapeDeserializer":
         if isinstance(source, bytes):
             source = BytesIO(source)
-        return _JSONShapeDeserializer(
-            source,
-            use_json_name=self._use_json_name,
-            use_timestamp_format=self._use_timestamp_format,
-            default_timestamp_format=self._default_timestamp_format,
-        )
+        return _JSONShapeDeserializer(source, settings=self._settings)

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

@@ -1,11 +1,21 @@
 #  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,6 +18,7 @@
 from smithy_core.traits import JSONNameTrait, TimestampFormatTrait
 from smithy_core.types import TimestampFormat
 
+from . import JSONSettings
 from .documents import JSONDocument
 
 # TODO: put these type hints in a pyi somewhere. There here because ijson isn't
@@ -89,17 +90,10 @@ def peek(self) -> JSONParseEvent:
 
 class JSONShapeDeserializer(ShapeDeserializer):
     def __init__(
-        self,
-        source: BytesReader,
-        *,
-        use_json_name: bool = True,
-        use_timestamp_format: bool = True,
-        default_timestamp_format: TimestampFormat = TimestampFormat.DATE_TIME,
+        self, source: BytesReader, *, settings: JSONSettings | None = None
     ) -> None:
         self._stream = BufferedParser(ijson.parse(source))
-        self._use_json_name = use_json_name
-        self._use_timestamp_format = use_timestamp_format
-        self._default_timestamp_format = default_timestamp_format
+        self._settings = settings or JSONSettings()
 
         # 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
@@ -164,13 +158,7 @@ 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,
-                use_json_name=self._use_json_name,
-                default_timestamp_format=self._default_timestamp_format,
-                use_timestamp_format=self._use_timestamp_format,
-            )
+            return JSONDocument(start.value, schema=schema, settings=self._settings)
 
         end_type = "end_map" if start.type == "start_map" else "end_array"
         builder = cast(TypedObjectBuilder, ObjectBuilder())
@@ -180,17 +168,11 @@ 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,
-            use_json_name=self._use_json_name,
-            default_timestamp_format=self._default_timestamp_format,
-            use_timestamp_format=self._use_timestamp_format,
-        )
+        return JSONDocument(builder.value, schema=schema, settings=self._settings)
 
     def read_timestamp(self, schema: Schema) -> datetime.datetime:
-        format = self._default_timestamp_format
-        if self._use_timestamp_format:
+        format = self._settings.default_timestamp_format
+        if self._settings.use_timestamp_format:
             if format_trait := schema.get_trait(TimestampFormatTrait):
                 format = format_trait.format
 
@@ -221,7 +203,7 @@ def read_struct(
         next(self._stream)
 
     def _resolve_member(self, schema: Schema, key: str) -> Schema | None:
-        if self._use_json_name:
+        if self._settings.use_json_name:
             if schema.id not in self._json_names:
                 self._cache_json_names(schema=schema)
             if key in self._json_names[schema.id]:

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

@@ -9,11 +9,12 @@
 from smithy_core.documents import Document, DocumentValue
 from smithy_core.prelude import DOCUMENT
 from smithy_core.schemas import Schema
-from smithy_core.shapes import ShapeType
+from smithy_core.shapes import ShapeID, ShapeType
 from smithy_core.traits import JSONNameTrait, TimestampFormatTrait
-from smithy_core.types import TimestampFormat
 from smithy_core.utils import expect_type
 
+from . import JSONSettings
+
 
 class JSONDocument(Document):
     _schema: Schema
@@ -24,24 +25,26 @@ def __init__(
         value: DocumentValue | dict[str, "Document"] | list["Document"],
         *,
         schema: Schema = DOCUMENT,
-        use_json_name: bool = True,
-        use_timestamp_format: bool = True,
-        default_timestamp_format: TimestampFormat = TimestampFormat.DATE_TIME,
+        settings: JSONSettings | None = None,
     ) -> None:
         super().__init__(value, schema=schema)
-        self._use_json_name = use_json_name
-        self._use_timestamp_format = use_timestamp_format
-        self._default_timestamp_format = default_timestamp_format
+        self._settings = settings or JSONSettings()
         self._json_names = {}
 
-        if use_json_name and schema.shape_type in (
+        if self._settings.use_json_name and schema.shape_type in (
             ShapeType.STRUCTURE,
             ShapeType.UNION,
         ):
             for member_name, member_schema in schema.members.items():
                 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))
 
@@ -51,8 +54,8 @@ def as_float(self) -> float:
         return float(expect_type(Decimal, self._value))
 
     def as_timestamp(self) -> datetime:
-        format = self._default_timestamp_format
-        if self._use_timestamp_format:
+        format = self._settings.default_timestamp_format
+        if self._settings.use_timestamp_format:
             if format_trait := self._schema.get_trait(TimestampFormatTrait):
                 format = format_trait.format
 
@@ -106,13 +109,7 @@ def _new_document(
         value: DocumentValue | dict[str, "Document"] | list["Document"],
         schema: Schema,
     ) -> "Document":
-        return JSONDocument(
-            value,
-            schema=schema,
-            use_json_name=self._use_json_name,
-            use_timestamp_format=self._use_timestamp_format,
-            default_timestamp_format=self._default_timestamp_format,
-        )
+        return JSONDocument(value, schema=schema, settings=self._settings)
 
     def _wrap_map(self, value: Mapping[str, DocumentValue]) -> dict[str, "Document"]:
         if self._schema.shape_type not in (ShapeType.STRUCTURE, ShapeType.UNION):

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

@@ -21,34 +21,23 @@
 from smithy_core.traits import JSONNameTrait, TimestampFormatTrait
 from smithy_core.types import TimestampFormat
 
-from . import Flushable
+from . import Flushable, JSONSettings
 
 _INF: float = float("inf")
 _NEG_INF: float = float("-inf")
 
 
 class JSONShapeSerializer(ShapeSerializer):
-    _stream: "StreamingJSONEncoder"
-    _use_json_name: bool
-    _use_timestamp_format: bool
-    _default_timestamp_format: TimestampFormat
-
     def __init__(
-        self,
-        sink: BytesWriter,
-        use_json_name: bool = True,
-        use_timestamp_format: bool = True,
-        default_timestamp_format: TimestampFormat = TimestampFormat.DATE_TIME,
+        self, sink: BytesWriter, *, settings: JSONSettings | None = None
     ) -> None:
         self._stream = StreamingJSONEncoder(sink)
-        self._use_json_name = use_json_name
-        self._use_timestamp_format = use_timestamp_format
-        self._default_timestamp_format = default_timestamp_format
+        self._settings = settings or JSONSettings()
 
     def begin_struct(
         self, schema: "Schema"
     ) -> AbstractContextManager["ShapeSerializer"]:
-        return JSONStructSerializer(self._stream, self, self._use_json_name)
+        return JSONStructSerializer(self._stream, self, self._settings.use_json_name)
 
     def begin_list(
         self, schema: "Schema", size: int
@@ -82,8 +71,8 @@ def write_blob(self, schema: "Schema", value: bytes) -> None:
         self._stream.write_string(b64encode(value).decode("utf-8"))
 
     def write_timestamp(self, schema: "Schema", value: datetime) -> None:
-        format = self._default_timestamp_format
-        if self._use_timestamp_format:
+        format = self._settings.default_timestamp_format
+        if self._settings.use_timestamp_format:
             if format_trait := schema.get_trait(TimestampFormatTrait):
                 format = format_trait.format