Address comments

Author: haydenbaker

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

@@ -640,22 +640,53 @@ 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:
+    """A registry for on-demand deserialization of types by using a mapping of shape IDs
+    to their deserializers."""
+
     def __init__(
         self,
         types: dict[ShapeID, type[DeserializeableShape]],
         sub_registry: "TypeRegistry | None" = None,
     ):
+        """Initialize a TypeRegistry.
+
+        :param types: A mapping of ShapeID to the shapes they deserialize to.
+        :param sub_registry: A registry to delegate to if an ID is not found in types.
+        """
         self._types = types
         self._sub_registry = sub_registry
 
     def get(self, shape: ShapeID) -> type[DeserializeableShape]:
+        """Get the deserializable shape for the given shape ID.
+
+        :param shape: The shape ID to get from the registry.
+        :returns: The corresponding deserializable shape.
+        :raises KeyError: If the shape ID is not found in the registry.
+        """
         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 __getitem__(self, shape: ShapeID):
+        """Get the deserializable shape for the given shape ID.
+
+        :param shape: The shape ID to get from the registry.
+        :returns: The corresponding deserializable shape.
+        :raises KeyError: If the shape ID is not found in the registry.
+        """
+        return self.get(shape)
+
+    def __contains__(self, shape: ShapeID):
+        """Check if the registry contains the given shape.
+
+        :param shape: The shape ID to check for.
+        """
+        return shape in self._types or (
+            self._sub_registry is not None and shape in self._sub_registry
+        )
+
     def deserialize(self, document: Document) -> DeserializeableShape:
         return document.as_shape(self.get(document.discriminator))

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

@@ -8,25 +8,44 @@
 def test_get():
     registry = TypeRegistry({ShapeID("com.example#Test"): TestShape})
 
-    result = registry.get(ShapeID("com.example#Test"))
+    result = registry[ShapeID("com.example#Test")]
 
     assert result == TestShape
 
 
+def test_contains():
+    registry = TypeRegistry({ShapeID("com.example#Test"): TestShape})
+
+    assert ShapeID("com.example#Test") in registry
+
+
 def test_get_sub_registry():
     sub_registry = TypeRegistry({ShapeID("com.example#Test"): TestShape})
     registry = TypeRegistry({}, sub_registry)
 
-    result = registry.get(ShapeID("com.example#Test"))
+    result = registry[ShapeID("com.example#Test")]
 
     assert result == TestShape
 
 
+def test_contains_sub_registry():
+    sub_registry = TypeRegistry({ShapeID("com.example#Test"): TestShape})
+    registry = TypeRegistry({}, sub_registry)
+
+    assert ShapeID("com.example#Test") in registry
+
+
 def test_get_no_match():
     registry = TypeRegistry({ShapeID("com.example#Test"): TestShape})
 
     with pytest.raises(KeyError, match="Unknown shape: com.example#Test2"):
-        registry.get(ShapeID("com.example#Test2"))
+        registry[ShapeID("com.example#Test2")]
+
+
+def test_contains_no_match():
+    registry = TypeRegistry({ShapeID("com.example#Test"): TestShape})
+
+    assert ShapeID("com.example#Test2") not in registry
 
 
 def test_deserialize():
@@ -39,6 +58,7 @@ def test_deserialize():
 
 
 class TestShape(DeserializeableShape):
+    __test__ = False
     schema = Schema(id=ShapeID("com.example#Test"), shape_type=ShapeType.STRING)
 
     def __init__(self, value: str):