Add a type registry

Author: haydenbaker

packages/smithy-core/src/smithy_core/aio/interfaces/__init__.py

@@ -1,11 +1,11 @@
 #  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 #  SPDX-License-Identifier: Apache-2.0
 from collections.abc import AsyncIterable
-from typing import Protocol, runtime_checkable, TYPE_CHECKING, Any
+from typing import Protocol, runtime_checkable, TYPE_CHECKING
 
 from ...interfaces import URI, Endpoint, TypedProperties
 from ...interfaces import StreamingBlob as SyncStreamingBlob
-
+from ...type_registry import TypeRegistry
 
 if TYPE_CHECKING:
     from ...schemas import APIOperation
@@ -126,7 +126,7 @@ async def deserialize_response[
         operation: "APIOperation[OperationInput, OperationOutput]",
         request: I,
         response: O,
-        error_registry: Any,  # TODO: add error registry
+        error_registry: TypeRegistry,
         context: TypedProperties,
     ) -> OperationOutput:
         """Deserializes the output from the tranport response or throws an exception.

packages/smithy-core/src/smithy_core/documents.py

@@ -143,6 +143,11 @@ def shape_type(self) -> ShapeType:
         """The Smithy data model type for the underlying contents of the document."""
         return self._type
 
+    @property
+    def discriminator(self) -> ShapeID:
+        """The shape ID that corresponds to the contents of the document."""
+        return self._schema.id
+
     def is_none(self) -> bool:
         """Indicates whether the document contains a null value."""
         return self._value is None and self._raw_value is None

packages/smithy-core/src/smithy_core/schemas.py

@@ -7,7 +7,7 @@
 from .exceptions import ExpectationNotMetException, SmithyException
 from .shapes import ShapeID, ShapeType
 from .traits import Trait, DynamicTrait, IdempotencyTokenTrait, StreamingTrait
-
+from .type_registry import TypeRegistry
 
 if TYPE_CHECKING:
     from .serializers import SerializeableShape
@@ -289,8 +289,7 @@ class APIOperation[I: "SerializeableShape", O: "DeserializeableShape"]:
     output_schema: Schema
     """The schema of the operation's output shape."""
 
-    # TODO: Add a type registry for errors
-    error_registry: Any
+    error_registry: TypeRegistry
     """A TypeRegistry used to create errors."""
 
     effective_auth_schemes: Sequence[ShapeID]

packages/smithy-core/src/smithy_core/type_registry.py

@@ -0,0 +1,30 @@
+#  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#  SPDX-License-Identifier: Apache-2.0
+
+from smithy_core.deserializers import (
+    DeserializeableShape,
+)  # TODO: fix typo in deserializable
+from smithy_core.documents import Document
+from smithy_core.shapes import ShapeID
+
+
+# A registry for on-demand deserialization of types by using a mapping of shape IDs to their deserializers.
+# TODO: protocol? Also, move into documents.py?
+class TypeRegistry:
+    def __init__(
+        self,
+        types: dict[ShapeID, type[DeserializeableShape]],
+        sub_registry: "TypeRegistry | None" = None,
+    ):
+        self._types = types
+        self._sub_registry = sub_registry
+
+    def get(self, shape: ShapeID) -> type[DeserializeableShape]:
+        if shape in self._types:
+            return self._types[shape]
+        if self._sub_registry is not None:
+            return self._sub_registry.get(shape)
+        raise KeyError(f"Unknown shape: {shape}")
+
+    def deserialize(self, document: Document) -> DeserializeableShape:
+        return document.as_shape(self.get(document.discriminator))

packages/smithy-core/tests/unit/test_type_registry.py

@@ -0,0 +1,48 @@
+from smithy_core.deserializers import DeserializeableShape, ShapeDeserializer
+from smithy_core.documents import Document
+from smithy_core.schemas import Schema
+from smithy_core.shapes import ShapeID, ShapeType
+from smithy_core.type_registry import TypeRegistry
+import pytest
+
+
+class TestTypeRegistry:
+    def test_get(self):
+        registry = TypeRegistry({ShapeID("com.example#Test"): TestShape})
+
+        result = registry.get(ShapeID("com.example#Test"))
+
+        assert result == TestShape
+
+    def test_get_sub_registry(self):
+        sub_registry = TypeRegistry({ShapeID("com.example#Test"): TestShape})
+        registry = TypeRegistry({}, sub_registry)
+
+        result = registry.get(ShapeID("com.example#Test"))
+
+        assert result == TestShape
+
+    def test_get_no_match(self):
+        registry = TypeRegistry({ShapeID("com.example#Test"): TestShape})
+
+        with pytest.raises(KeyError, match="Unknown shape: com.example#Test2"):
+            registry.get(ShapeID("com.example#Test2"))
+
+    def test_deserialize(self):
+        shape_id = ShapeID("com.example#Test")
+        registry = TypeRegistry({shape_id: TestShape})
+
+        result = registry.deserialize(Document("abc123", schema=TestShape.schema))
+
+        assert isinstance(result, TestShape) and result.value == "abc123"
+
+
+class TestShape(DeserializeableShape):
+    schema = Schema(id=ShapeID("com.example#Test"), shape_type=ShapeType.STRING)
+
+    def __init__(self, value: str):
+        self.value = value
+
+    @classmethod
+    def deserialize(cls, deserializer: ShapeDeserializer) -> "TestShape":
+        return TestShape(deserializer.read_string(schema=TestShape.schema))