Typed context

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

@@ -3,7 +3,7 @@
 from collections.abc import AsyncIterable
 from typing import Protocol, runtime_checkable, TYPE_CHECKING, Any
 
-from ...interfaces import URI, Endpoint
+from ...interfaces import URI, Endpoint, TypedProperties
 from ...interfaces import StreamingBlob as SyncStreamingBlob
 
 
@@ -93,7 +93,7 @@ def serialize_request[
         operation: "APIOperation[OperationInput, OperationOutput]",
         input: OperationInput,
         endpoint: URI,
-        context: dict[str, Any],
+        context: TypedProperties,
     ) -> I:
         """Serialize an operation input into a transport request.
 
@@ -127,7 +127,7 @@ async def deserialize_response[
         request: I,
         response: O,
         error_registry: Any,  # TODO: add error registry
-        context: dict[str, Any],  # TODO: replace with a typed context bag
+        context: TypedProperties,
     ) -> OperationOutput:
         """Deserializes the output from the tranport response or throws an exception.
 

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

@@ -1,7 +1,9 @@
 #  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 #  SPDX-License-Identifier: Apache-2.0
 from copy import copy, deepcopy
-from typing import Any, TypeVar
+from typing import TypeVar
+
+from .types import TypedProperties
 
 Request = TypeVar("Request")
 Response = TypeVar("Response")
@@ -34,7 +36,7 @@ def __init__(
         self._response = response
         self._transport_request = transport_request
         self._transport_response = transport_response
-        self._properties: dict[str, Any] = {}
+        self._properties = TypedProperties()
 
     @property
     def request(self) -> Request:
@@ -73,7 +75,7 @@ def transport_response(self) -> TransportResponse:
         return self._transport_response
 
     @property
-    def properties(self) -> dict[str, Any]:
+    def properties(self) -> TypedProperties:
         """Retrieve the generic property bag.
 
         These untyped properties will be made available to all other interceptors or

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

@@ -1,7 +1,17 @@
 # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 # SPDX-License-Identifier: Apache-2.0
 from asyncio import iscoroutinefunction
-from typing import Protocol, runtime_checkable, Any, TypeGuard
+from typing import (
+    Protocol,
+    runtime_checkable,
+    Any,
+    TypeGuard,
+    overload,
+    Iterator,
+    KeysView,
+    ValuesView,
+    ItemsView,
+)
 
 
 class URI(Protocol):
@@ -99,3 +109,91 @@ class Endpoint(Protocol):
     For example, in some AWS use cases this might contain HTTP headers to add to each
     request.
     """
+
+
+@runtime_checkable
+class PropertyKey[T](Protocol):
+    """A typed properties key.
+
+    Used with :py:class:`Context` to set and get typed values.
+
+    For a concrete implementation, see :py:class:`smithy_core.types.PropertyKey`.
+    """
+
+    key: str
+    """The string key used to access the value."""
+
+    value_type: type[T]
+    """The type of the associated value in the properties bag."""
+
+    def __str__(self) -> str:
+        return self.key
+
+
+# This is currently strongly tied to being compatible with a dict[str, Any], but we
+# could remove that to allow for potentially more efficient maps. That might introduce
+# unacceptable usability penalties or footguns though.
+@runtime_checkable
+class TypedProperties(Protocol):
+    """A properties map with typed setters and getters.
+
+    Keys can be either a string or a :py:class:`PropertyKey`. Using a PropertyKey instead
+    of a string enables type checkers to narrow to the associated value type rather
+    than having to use Any.
+
+    Letting the value be either a string or PropertyKey allows consumers who care about
+    typing to get it, and those who don't care about typing to not have to think about
+    it.
+
+    ..code-block:: python
+
+        foo = PropertyKey(key="foo", value_type=str)
+        properties = TypedProperties()
+        properties[foo] = "bar"
+
+        assert assert_type(properties[foo], str) == "bar
+        assert assert_type(properties["foo"], Any) == "bar
+
+
+    For a concrete implementation, see :py:class:`smithy_core.types.TypedProperties`.
+    """
+
+    @overload
+    def __getitem__[T](self, key: PropertyKey[T]) -> T: ...
+    @overload
+    def __getitem__(self, key: str) -> Any: ...
+
+    @overload
+    def __setitem__[T](self, key: PropertyKey[T], value: T) -> None: ...
+    @overload
+    def __setitem__(self, key: str, value: Any) -> None: ...
+
+    def __delitem__(self, key: str | PropertyKey[Any]) -> None: ...
+
+    @overload
+    def get[T](self, key: PropertyKey[T], default: None = None) -> T | None: ...
+    @overload
+    def get[T](self, key: PropertyKey[T], default: T) -> T: ...
+    @overload
+    def get[T, DT](self, key: PropertyKey[T], default: DT) -> T | DT: ...
+    @overload
+    def get(self, key: str, default: None = None) -> Any | None: ...
+    @overload
+    def get[T](self, key: str, default: T) -> Any | T: ...
+
+    @overload
+    def pop[T](self, key: PropertyKey[T], default: None = None) -> T | None: ...
+    @overload
+    def pop[T](self, key: PropertyKey[T], default: T) -> T: ...
+    @overload
+    def pop[T, DT](self, key: PropertyKey[T], default: DT) -> T | DT: ...
+    @overload
+    def pop(self, key: str, default: None = None) -> Any | None: ...
+    @overload
+    def pop[T](self, key: str, default: T) -> Any | T: ...
+
+    def __iter__(self) -> Iterator[str]: ...
+    def items(self) -> ItemsView[str, Any]: ...
+    def keys(self) -> KeysView[str]: ...
+    def values(self) -> ValuesView[Any]: ...
+    def __contains__(self, key: object) -> bool: ...

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

@@ -2,11 +2,13 @@
 #  SPDX-License-Identifier: Apache-2.0
 import json
 import re
+import sys
+from collections import UserDict
 from collections.abc import Mapping, Sequence
 from datetime import datetime
 from email.utils import format_datetime, parsedate_to_datetime
 from enum import Enum
-from typing import Any
+from typing import Any, overload
 from dataclasses import dataclass
 
 from .exceptions import ExpectationNotMetException
@@ -17,6 +19,8 @@
     serialize_epoch_seconds,
     serialize_rfc3339,
 )
+from .interfaces import PropertyKey as _PropertyKey
+from .interfaces import TypedProperties as _TypedProperties
 
 _GREEDY_LABEL_RE = re.compile(r"\{(\w+)\+\}")
 
@@ -153,3 +157,99 @@ def format(self, *args: object, **kwargs: str) -> str:
                 f'Path must not contain empty segments, but was "{result}".'
             )
         return result
+
+
+@dataclass(kw_only=True, frozen=True, slots=True, init=False)
+class PropertyKey[T](_PropertyKey[T]):
+    """A typed property key."""
+
+    key: str
+    """The string key used to access the value."""
+
+    value_type: type[T]
+    """The type of the associated value in the property bag."""
+
+    def __init__(self, *, key: str, value_type: type[T]) -> None:
+        # Intern the key to speed up dict access
+        object.__setattr__(self, "key", sys.intern(key))
+        object.__setattr__(self, "value_type", value_type)
+
+
+class TypedProperties(UserDict[str, Any], _TypedProperties):
+    """A map with typed setters and getters.
+
+    Keys can be either a string or a :py:class:`smithy_core.interfaces.PropertyKey`.
+    Using a PropertyKey instead of a string enables type checkers to narrow to the
+    associated value type rather than having to use Any.
+
+    Letting the value be either a string or PropertyKey allows consumers who care about
+    typing to get it, and those who don't care about typing to not have to think about
+    it.
+
+    ..code-block:: python
+
+        foo = PropertyKey(key="foo", value_type=str)
+        properties = TypedProperties()
+        properties[foo] = "bar"
+
+        assert assert_type(properties[foo], str) == "bar
+        assert assert_type(properties["foo"], Any) == "bar
+    """
+
+    @overload
+    def __getitem__[T](self, key: _PropertyKey[T]) -> T: ...
+    @overload
+    def __getitem__(self, key: str) -> Any: ...
+    def __getitem__(self, key: str | _PropertyKey[Any]) -> Any:
+        return self.data[key if isinstance(key, str) else key.key]
+
+    @overload
+    def __setitem__[T](self, key: _PropertyKey[T], value: T) -> None: ...
+    @overload
+    def __setitem__(self, key: str, value: Any) -> None: ...
+    def __setitem__(self, key: str | _PropertyKey[Any], value: Any) -> None:
+        if isinstance(key, _PropertyKey):
+            if not isinstance(value, key.value_type):
+                raise ValueError(
+                    f"Expected value type of {key.value_type}, but was {type(value)}"
+                )
+            key = key.key
+        self.data[key] = value
+
+    def __delitem__(self, key: str | _PropertyKey[Any]) -> None:
+        del self.data[key if isinstance(key, str) else key.key]
+
+    def __contains__(self, key: object) -> bool:
+        return super().__contains__(key.key if isinstance(key, _PropertyKey) else key)
+
+    @overload
+    def get[T](self, key: _PropertyKey[T], default: None = None) -> T | None: ...
+    @overload
+    def get[T](self, key: _PropertyKey[T], default: T) -> T: ...
+    @overload
+    def get[T, DT](self, key: _PropertyKey[T], default: DT) -> T | DT: ...
+    @overload
+    def get(self, key: str, default: None = None) -> Any | None: ...
+    @overload
+    def get[T](self, key: str, default: T) -> Any | T: ...
+
+    # pyright has trouble detecting compatible overrides when both the superclass
+    # and subclass have overloads.
+    def get(self, key: str | _PropertyKey[Any], default: Any = None) -> Any:  # type: ignore
+        return self.data.get(key if isinstance(key, str) else key.key, default)
+
+    @overload
+    def pop[T](self, key: _PropertyKey[T], default: None = None) -> T | None: ...
+    @overload
+    def pop[T](self, key: _PropertyKey[T], default: T) -> T: ...
+    @overload
+    def pop[T, DT](self, key: _PropertyKey[T], default: DT) -> T | DT: ...
+    @overload
+    def pop(self, key: str, default: None = None) -> Any | None: ...
+    @overload
+    def pop[T](self, key: str, default: T) -> Any | T: ...
+
+    # pyright has trouble detecting compatible overrides when both the superclass
+    # and subclass have overloads.
+    def pop(self, key: str | _PropertyKey[Any], default: Any = None) -> Any:  # type: ignore
+        return self.data.pop(key if isinstance(key, str) else key.key, default)