Move type registry into documents, fix circular import

Author: haydenbaker

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

@@ -5,7 +5,7 @@
 
 from ...interfaces import URI, Endpoint, TypedProperties
 from ...interfaces import StreamingBlob as SyncStreamingBlob
-from ...type_registry import TypeRegistry
+from ...documents import TypeRegistry
 
 if TYPE_CHECKING:
     from ...schemas import APIOperation

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

@@ -638,3 +638,24 @@ def read_document(self, schema: "Schema") -> Document:
     @override
     def read_timestamp(self, schema: "Schema") -> datetime.datetime:
         return self._value.as_timestamp()
+
+
+# A registry for on-demand deserialization of types by using a mapping of shape IDs to their deserializers.
+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/src/smithy_core/schemas.py

@@ -7,9 +7,9 @@
 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 .documents import TypeRegistry
     from .serializers import SerializeableShape
     from .deserializers import DeserializeableShape
 
@@ -289,7 +289,7 @@ class APIOperation[I: "SerializeableShape", O: "DeserializeableShape"]:
     output_schema: Schema
     """The schema of the operation's output shape."""
 
-    error_registry: TypeRegistry
+    error_registry: "TypeRegistry"
     """A TypeRegistry used to create errors."""
 
     effective_auth_schemes: Sequence[ShapeID]

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

@@ -1,40 +1,41 @@
 from smithy_core.deserializers import DeserializeableShape, ShapeDeserializer
-from smithy_core.documents import Document
+from smithy_core.documents import Document, TypeRegistry
 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})
+def test_get():
+    registry = TypeRegistry({ShapeID("com.example#Test"): TestShape})
 
-        result = registry.get(ShapeID("com.example#Test"))
+    result = registry.get(ShapeID("com.example#Test"))
 
-        assert result == TestShape
+    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"))
+def test_get_sub_registry():
+    sub_registry = TypeRegistry({ShapeID("com.example#Test"): TestShape})
+    registry = TypeRegistry({}, sub_registry)
 
-        assert result == TestShape
+    result = registry.get(ShapeID("com.example#Test"))
 
-    def test_get_no_match(self):
-        registry = TypeRegistry({ShapeID("com.example#Test"): TestShape})
+    assert result == 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})
+def test_get_no_match():
+    registry = TypeRegistry({ShapeID("com.example#Test"): TestShape})
 
-        result = registry.deserialize(Document("abc123", schema=TestShape.schema))
+    with pytest.raises(KeyError, match="Unknown shape: com.example#Test2"):
+        registry.get(ShapeID("com.example#Test2"))
 
-        assert isinstance(result, TestShape) and result.value == "abc123"
+
+def test_deserialize():
+    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):