Implement __contains__ for Schema

Author: JordonPhillips

packages/aws-event-stream/src/aws_event_stream/_private/__init__.py

@@ -10,6 +10,6 @@
 
 def get_payload_member(schema: Schema) -> Schema | None:
     for member in schema.members.values():
-        if EventPayloadTrait.id in member.traits:
+        if EventPayloadTrait in member:
             return member
     return None

packages/aws-event-stream/src/aws_event_stream/_private/deserializers.py

@@ -158,10 +158,7 @@ def read_struct(
         headers_deserializer = EventHeaderDeserializer(self._headers)
         for key in self._headers.keys():
             member_schema = schema.members.get(key)
-            if (
-                member_schema is not None
-                and EventHeaderTrait.id in member_schema.traits
-            ):
+            if member_schema is not None and EventHeaderTrait in member_schema:
                 consumer(member_schema, headers_deserializer)
 
         if self._payload_deserializer:

packages/aws-event-stream/src/aws_event_stream/_private/serializers.py

@@ -102,7 +102,7 @@ def begin_struct(self, schema: "Schema") -> Iterator[ShapeSerializer]:
 
         headers_encoder = EventHeaderEncoder()
 
-        if ErrorTrait.id in schema.traits:
+        if ErrorTrait in schema:
             headers_encoder.encode_string(":message-type", "exception")
             headers_encoder.encode_string(
                 ":exception-type", schema.expect_member_name()
@@ -214,7 +214,7 @@ def __init__(
         self._payload_struct_serializer = payload_struct_serializer
 
     def before(self, schema: "Schema") -> ShapeSerializer:
-        if EventHeaderTrait.id in schema.traits:
+        if EventHeaderTrait in schema:
             return self._header_serializer
         return self._payload_struct_serializer
 

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

@@ -2,13 +2,11 @@
 #  SPDX-License-Identifier: Apache-2.0
 from collections.abc import Mapping
 from dataclasses import dataclass, field, replace
-from typing import TYPE_CHECKING, NotRequired, Required, Self, TypedDict, overload
+from typing import NotRequired, Required, Self, TypedDict, overload, Any
 
 from .exceptions import ExpectationNotMetException, SmithyException
 from .shapes import ShapeID, ShapeType
-
-if TYPE_CHECKING:
-    from .traits import Trait, DynamicTrait
+from .traits import Trait, DynamicTrait
 
 
 @dataclass(kw_only=True, frozen=True, init=False)
@@ -164,6 +162,24 @@ def expect_trait(self, t: "type[Trait] | ShapeID") -> "Trait | DynamicTrait":
         id = t if isinstance(t, ShapeID) else t.id
         return self.traits[id]
 
+    def __contains__(self, item: Any):
+        """Returns whether the schema has the given member or trait."""
+        match item:
+            case type():
+                if issubclass(item, Trait):
+                    return item.id in self.traits
+                return False
+            case ShapeID():
+                if (member := item.member) is not None:
+                    if self.id.with_member(member) == item:
+                        return member in self.members
+                    return False
+                return item in self.traits
+            case str():
+                return item in self.members
+            case _:
+                return False
+
     @classmethod
     def collection(
         cls,

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

@@ -2,10 +2,16 @@
 
 import pytest
 
+from typing import Any
+
 from smithy_core.exceptions import ExpectationNotMetException
 from smithy_core.schemas import Schema
 from smithy_core.shapes import ShapeID, ShapeType
-from smithy_core.traits import InternalTrait, DynamicTrait, SensitiveTrait
+from smithy_core.traits import (
+    InternalTrait,
+    DynamicTrait,
+    SensitiveTrait,
+)
 
 ID: ShapeID = ShapeID("ns.foo#bar")
 STRING = Schema(id=ShapeID("smithy.api#String"), shape_type=ShapeType.STRING)
@@ -143,3 +149,25 @@ def test_member_constructor_asserts_target_is_not_member():
     )
     with pytest.raises(ExpectationNotMetException):
         Schema.member(id=ShapeID("smithy.example#Foo$bar"), target=target, index=0)
+
+
+@pytest.mark.parametrize(
+    "item, contains",
+    [
+        (SensitiveTrait, True),
+        (SensitiveTrait.id, True),
+        (InternalTrait, False),
+        (InternalTrait.id, False),
+        ("baz", True),
+        (ID.with_member("baz"), True),
+        (ID, False),
+    ],
+)
+def test_contains(item: Any, contains: bool):
+    schema = Schema.collection(
+        id=ID,
+        members={"baz": {"target": STRING, "index": 0}},
+        traits=[SensitiveTrait()],
+    )
+
+    assert (item in schema) == contains