Update traits to allow for dynamically discoverable implementations

codegen/core/src/main/java/software/amazon/smithy/python/codegen/generators/SchemaGenerator.java

@@ -122,7 +122,7 @@ private void writeTraits(PythonWriter writer, Map<ShapeId, Optional<Node>> trait
         writer.putContext("traits", traits);
         writer.write("""
                 ${#traits}
-                Trait(id=ShapeID(${key:S})${?value}, value=${value:N}${/value}),
+                Trait.new(id=ShapeID(${key:S})${?value}, value=${value:N}${/value}),
                 ${/traits}""");
         writer.popState();
     }

designs/serialization.md

@@ -80,27 +80,28 @@ implementation and/or additional helper methods.
 class Schema:
     id: ShapeID
     shape_type: ShapeType
-    traits: dict[ShapeID, "Trait"] = field(default_factory=dict)
+    traits: dict[ShapeID, "Trait | DynamicTrait"] = field(default_factory=dict)
     members: dict[str, "Schema"] = field(default_factory=dict)
     member_target: "Schema | None" = None
     member_index: int | None = None
 
+    @overload
+    def get_trait[T: "Trait"](self, t: type[T]) -> T | None: ...
+    @overload
+    def get_trait(self, t: ShapeID) -> "Trait | DynamicTrait | None": ...
+    def get_trait(self, t: "type[Trait] | ShapeID") -> "Trait | DynamicTrait | None":\
+        return self.traits.get(t if isinstance(t, ShapeID) else t.id)
+
     @classmethod
     def collection(
         cls,
         *,
         id: ShapeID,
         shape_type: ShapeType = ShapeType.STRUCTURE,
-        traits: list["Trait"] | None = None,
+        traits: list["Trait | DynamicTrait"] | None = None,
         members: Mapping[str, "MemberSchema"] | None = None,
     ) -> Self:
         ...
-
-
-@dataclass(kw_only=True, frozen=True)
-class Trait:
-    id: "ShapeID"
-    value: "DocumentValue" = field(default_factory=dict)
 ```
 
 Below is an example Smithy `structure` shape, followed by the `Schema` it would
@@ -122,13 +123,112 @@ EXAMPLE_STRUCTURE_SCHEMA = Schema.collection(
             "target": INTEGER,
             "index": 0,
             "traits": [
-                Trait(id=ShapeID("smithy.api#default"), value=0),
+                DefaultTrait(0),
             ],
         },
     },
 )
 ```
 
+### Traits
+
+Traits are model components that can be attached to shapes to describe
+additional information about the shape; shapes provide the structure and layout
+of an API, while traits provide refinement and style. Smithy provides a number
+of built-in traits, plus a number of additional traits that may be found in
+first-party dependencies. In addition to those first-party traits, traits may be
+defined externally.
+
+In Python, there are two kinds of traits. The first is the `DynamicTrait`. This
+represents traits that have no known associated Python class. Traits not defined
+by Smithy itself may be unknown, for example, but still need representation.
+
+The other kind of trait inherits from the `Trait` class. This represents known
+traits, such as those defined by Smithy itself or those defined externally but
+made available in Python. Since these are concrete classes, they may be more
+comfortable to use, providing better typed accessors to data or even relevant
+utility functions.
+
+Both kinds of traits implement an inherent `Protocol` - they both have the `id`
+and `document_value` properties with identical type signatures. This allows them
+to be used interchangeably for those that don't care about the concrete types.
+It also allows concrete types to be introduced later without a breaking change.
+
+
+```python
+@dataclass(kw_only=True, frozen=True, slots=True)
+class DynamicTrait:
+    id: ShapeID
+    document_value: DocumentValue = None
+
+
+@dataclass(init=False, frozen=True)
+class Trait:
+
+    _REGISTRY: ClassVar[dict[ShapeID, type["Trait"]]] = {}
+
+    id: ClassVar[ShapeID]
+
+    document_value: DocumentValue = None
+
+    def __init_subclass__(cls, id: ShapeID) -> None:
+        cls.id = id
+        Trait._REGISTRY[id] = cls
+
+    def __init__(self, value: DocumentValue | DynamicTrait = None):
+        if type(self) is Trait:
+            raise TypeError(
+                "Only subclasses of Trait may be directly instantiated. "
+                "Use DynamicTrait for traits without a concrete class."
+            )
+
+        if isinstance(value, DynamicTrait):
+            if value.id != self.id:
+                raise ValueError(
+                    f"Attempted to instantiate an instance of {type(self)} from an "
+                    f"invalid ID. Expected {self.id} but found {value.id}."
+                )
+            # Note that setattr is needed because it's a frozen (read-only) dataclass
+            object.__setattr__(self, "document_value", value.document_value)
+        else:
+            object.__setattr__(self, "document_value", value)
+
+    # Dynamically creates a subclass instance based on the trait id
+    @staticmethod
+    def new(id: ShapeID, value: "DocumentValue" = None) -> "Trait | DynamicTrait":
+        if (cls := Trait._REGISTRY.get(id, None)) is not None:
+            return cls(value)
+        return DynamicTrait(id=id, document_value=value)
+```
+
+The `Trait` class implements a dynamic registry that allows it to know about
+trait implementations automatically. The base class maintains a mapping of trait
+ID to the trait class. Since implementations must all share the same constructor
+signature, it can then use that registry to dynamically construct concrete types
+it knows about in the `new` factory method with a fallback to `DynamicTrait`.
+
+The `new` factory method will be used to construct traits when `Schema`s are
+generated, so any generated schemas will be able to take advantage of the
+registry.
+
+Below is an example of a `Trait` implementation.
+
+```python
+@dataclass(init=False, frozen=True)
+class TimestampFormatTrait(Trait, id=ShapeID("smithy.api#timestampFormat")):
+    format: TimestampFormat
+
+    def __init__(self, value: "DocumentValue | DynamicTrait" = None):
+        super().__init__(value)
+        assert isinstance(self.document_value, str)
+        object.__setattr__(self, "format", TimestampFormat(self.document_value))
+```
+
+Data in traits is intended to be immutable, so both `DynamicTrait` and `Trait`
+are dataclasses with `frozen=True`, and all implementations of `Trait` must also
+use that argument. This can be worked around during `__init__` using
+`object.__setattr__` to set any additional properties the `Trait` defines.
+
 ## Shape Serializers and Serializeable Shapes
 
 Serialization will function by the interaction of two interfaces:

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

@@ -2,14 +2,14 @@
 # SPDX-License-Identifier: Apache-2.0
 from smithy_core.schemas import Schema
 
-from .traits import EVENT_PAYLOAD_TRAIT
+from smithy_core.traits import EventPayloadTrait
 
 INITIAL_REQUEST_EVENT_TYPE = "initial-request"
 INITIAL_RESPONSE_EVENT_TYPE = "initial-response"
 
 
 def get_payload_member(schema: Schema) -> Schema | None:
     for member in schema.members.values():
-        if EVENT_PAYLOAD_TRAIT in member.traits:
+        if EventPayloadTrait in member:
             return member
     return None

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

@@ -23,7 +23,7 @@
     INITIAL_RESPONSE_EVENT_TYPE,
     get_payload_member,
 )
-from .traits import EVENT_HEADER_TRAIT
+from smithy_core.traits import EventHeaderTrait
 
 INITIAL_MESSAGE_TYPES = (INITIAL_REQUEST_EVENT_TYPE, INITIAL_RESPONSE_EVENT_TYPE)
 
@@ -158,7 +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 EVENT_HEADER_TRAIT 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

@@ -18,7 +18,6 @@
     SpecificShapeSerializer,
 )
 from smithy_core.shapes import ShapeType
-from smithy_core.utils import expect_type
 from smithy_event_stream.aio.interfaces import AsyncEventPublisher
 
 from ..events import EventHeaderEncoder, EventMessage
@@ -28,7 +27,7 @@
     INITIAL_RESPONSE_EVENT_TYPE,
     get_payload_member,
 )
-from .traits import ERROR_TRAIT, EVENT_HEADER_TRAIT, MEDIA_TYPE_TRAIT
+from smithy_core.traits import ErrorTrait, EventHeaderTrait, MediaTypeTrait
 
 _DEFAULT_STRING_CONTENT_TYPE = "text/plain"
 _DEFAULT_BLOB_CONTENT_TYPE = "application/octet-stream"
@@ -103,7 +102,7 @@ def begin_struct(self, schema: "Schema") -> Iterator[ShapeSerializer]:
 
         headers_encoder = EventHeaderEncoder()
 
-        if ERROR_TRAIT in schema.traits:
+        if ErrorTrait in schema:
             headers_encoder.encode_string(":message-type", "exception")
             headers_encoder.encode_string(
                 ":exception-type", schema.expect_member_name()
@@ -146,8 +145,8 @@ def begin_struct(self, schema: "Schema") -> Iterator[ShapeSerializer]:
         )
 
     def _get_payload_media_type(self, schema: Schema, default: str) -> str:
-        if (media_type := schema.traits.get(MEDIA_TYPE_TRAIT)) is not None:
-            return expect_type(str, media_type.value)
+        if (media_type := schema.get_trait(MediaTypeTrait)) is not None:
+            return media_type.value
 
         match schema.shape_type:
             case ShapeType.STRING:
@@ -215,7 +214,7 @@ def __init__(
         self._payload_struct_serializer = payload_struct_serializer
 
     def before(self, schema: "Schema") -> ShapeSerializer:
-        if EVENT_HEADER_TRAIT in schema.traits:
+        if EventHeaderTrait in schema:
             return self._header_serializer
         return self._payload_struct_serializer
 

packages/aws-event-stream/tests/unit/_private/__init__.py

@@ -19,15 +19,21 @@
 from smithy_core.schemas import Schema
 from smithy_core.serializers import ShapeSerializer
 from smithy_core.shapes import ShapeID, ShapeType
-from smithy_core.traits import Trait
+from smithy_core.traits import (
+    EventHeaderTrait,
+    EventPayloadTrait,
+    ErrorTrait,
+    RequiredTrait,
+    StreamingTrait,
+)
 
 from aws_event_stream.events import Byte, EventMessage, Long, Short
 
-EVENT_HEADER_TRAIT = Trait(id=ShapeID("smithy.api#eventHeader"))
-EVENT_PAYLOAD_TRAIT = Trait(id=ShapeID("smithy.api#eventPayload"))
-ERROR_TRAIT = Trait(id=ShapeID("smithy.api#error"), value="client")
-REQUIRED_TRAIT = Trait(id=ShapeID("smithy.api#required"))
-STREAMING_TRAIT = Trait(id=ShapeID("smith.api#streaming"))
+EVENT_HEADER_TRAIT = EventHeaderTrait()
+EVENT_PAYLOAD_TRAIT = EventPayloadTrait()
+ERROR_TRAIT = ErrorTrait("client")
+REQUIRED_TRAIT = RequiredTrait()
+STREAMING_TRAIT = StreamingTrait()
 
 
 SCHEMA_MESSAGE_EVENT = Schema.collection(

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

@@ -4,7 +4,8 @@
 
 from .schemas import Schema
 from .shapes import ShapeID, ShapeType
-from .traits import Trait
+from .traits import DefaultTrait, UnitTypeTrait
+
 
 BLOB = Schema(
     id=ShapeID("smithy.api#Blob"),
@@ -71,54 +72,50 @@
     shape_type=ShapeType.DOCUMENT,
 )
 
-
-_DEFAULT = ShapeID("smithy.api#default")
-
-
 PRIMITIVE_BOOLEAN = Schema(
     id=ShapeID("smithy.api#PrimitiveBoolean"),
     shape_type=ShapeType.BOOLEAN,
-    traits=[Trait(id=_DEFAULT, value=False)],
+    traits=[DefaultTrait(False)],
 )
 
 PRIMITIVE_BYTE = Schema(
     id=ShapeID("smithy.api#PrimitiveByte"),
     shape_type=ShapeType.BYTE,
-    traits=[Trait(id=_DEFAULT, value=0)],
+    traits=[DefaultTrait(0)],
 )
 
 PRIMITIVE_SHORT = Schema(
     id=ShapeID("smithy.api#PrimitiveShort"),
     shape_type=ShapeType.SHORT,
-    traits=[Trait(id=_DEFAULT, value=0)],
+    traits=[DefaultTrait(0)],
 )
 
 PRIMITIVE_INTEGER = Schema(
     id=ShapeID("smithy.api#PrimitiveInteger"),
     shape_type=ShapeType.INTEGER,
-    traits=[Trait(id=_DEFAULT, value=0)],
+    traits=[DefaultTrait(0)],
 )
 
 PRIMITIVE_LONG = Schema(
     id=ShapeID("smithy.api#PrimitiveLong"),
     shape_type=ShapeType.LONG,
-    traits=[Trait(id=_DEFAULT, value=0)],
+    traits=[DefaultTrait(0)],
 )
 
 PRIMITIVE_FLOAT = Schema(
     id=ShapeID("smithy.api#PrimitiveFloat"),
     shape_type=ShapeType.FLOAT,
-    traits=[Trait(id=_DEFAULT, value=0.0)],
+    traits=[DefaultTrait(0)],
 )
 
 PRIMITIVE_DOUBLE = Schema(
     id=ShapeID("smithy.api#PrimitiveDouble"),
     shape_type=ShapeType.DOUBLE,
-    traits=[Trait(id=_DEFAULT, value=0.0)],
+    traits=[DefaultTrait(0)],
 )
 
 UNIT = Schema(
     id=ShapeID("smithy.api#Unit"),
-    shape_type=ShapeType.DOUBLE,
-    traits=[Trait(id=ShapeID("smithy.api#UnitTypeTrait"))],
+    shape_type=ShapeType.STRUCTURE,
+    traits=[UnitTypeTrait()],
 )