Typing prototype

halogen/appel.py

@@ -0,0 +1,3 @@
+from .test import PaginationResponse
+
+x = PaginationResponse.serialize({"test"})

halogen/schema.py

@@ -2,7 +2,18 @@
 
 import inspect
 from collections import OrderedDict, namedtuple
-from typing import Iterable, Optional, Union
+from typing import (
+    Any,
+    Iterable,
+    Optional,
+    Union,
+    TypeVar,
+    Generic,
+    overload,
+    TYPE_CHECKING,
+    TypedDict,
+    List as TypingList,
+)
 
 from cached_property import cached_property
 
@@ -59,7 +70,9 @@ def get(self, obj, **kwargs):
         if callable(self.getter):
             return self.getter(obj, **_get_context(self._getter_argspec, kwargs))
 
-        assert isinstance(self.getter, str), "Accessor must be a function or a dot-separated string."
+        assert isinstance(self.getter, str), (
+            "Accessor must be a function or a dot-separated string."
+        )
 
         if obj is None:
             return None
@@ -89,7 +102,9 @@ def set(self, obj, value):
         if callable(self.setter):
             return self.setter(obj, value)
 
-        assert isinstance(self.setter, str), "Accessor must be a function or a dot-separated string."
+        assert isinstance(self.setter, str), (
+            "Accessor must be a function or a dot-separated string."
+        )
 
         def _set(obj, attr, value):
             if isinstance(obj, dict):
@@ -106,15 +121,57 @@ def _set(obj, attr, value):
 
     def __repr__(self):
         """Accessor representation."""
-        return "<{0} getter='{1}', setter='{2}'>".format(self.__class__.__name__, self.getter, self.setter)
+        return "<{0} getter='{1}', setter='{2}'>".format(
+            self.__class__.__name__, self.getter, self.setter
+        )
 
 
-class Attr(object):
+T = TypeVar("T")
+
+
+class Attr(Generic[T]):
     """Schema attribute."""
 
     creation_counter = 0
 
-    def __init__(self, attr_type=None, attr=None, required: bool = True, exclude: Optional[Iterable] = None, **kwargs):
+    @overload
+    def __init__(
+        self,
+        attr_type: "types.Type[T]" = ...,
+        attr=None,
+        required: bool = True,
+        exclude: Optional[Iterable] = None,
+        **kwargs,
+    ) -> None: ...
+
+    @overload
+    def __init__(
+        self,
+        attr_type: type["_Schema"] = ...,
+        attr=None,
+        required: bool = True,
+        exclude: Optional[Iterable] = None,
+        **kwargs,
+    ) -> None: ...
+
+    @overload
+    def __init__(
+        self,
+        attr_type: object = ...,
+        attr=None,
+        required: bool = True,
+        exclude: Optional[Iterable] = None,
+        **kwargs,
+    ) -> None: ...
+
+    def __init__(
+        self,
+        attr_type=None,
+        attr=None,
+        required: bool = True,
+        exclude: Optional[Iterable] = None,
+        **kwargs,
+    ):
         """Attribute constructor.
 
         :param attr_type: Type, Schema or constant that does the type conversion of the attribute.
@@ -132,6 +189,20 @@ def __init__(self, attr_type=None, attr=None, required: bool = True, exclude: Op
         self.creation_counter = Attr.creation_counter
         Attr.creation_counter += 1
 
+    @overload
+    def __get__(self, obj: None, objtype: Optional[type] = None) -> "Attr[T]": ...
+
+    @overload
+    def __get__(self, obj: object, objtype: Optional[type] = None) -> "Attr[T]": ...
+
+    def __get__(self, obj: Optional[object], objtype: Optional[type] = None):
+        # Attr instances are removed from Schema classes at runtime, so this is a typing-only aid.
+        return self
+
+    def __set__(self, obj: object, value: T) -> None:
+        # Typing-only descriptor hook.
+        raise AttributeError("Attr descriptors are not set on instances.")
+
     @property
     def compartment(self):
         """The key of the compartment this attribute will be placed into (for example: _links or _embedded)."""
@@ -184,8 +255,12 @@ def serialize(self, value, **kwargs):
                     raise
                 value = self._default()
 
-            value = self.attr_type.serialize(value, **_get_context(self._attr_type_serialize_argspec, kwargs))
-            value = self._default() if value is None and hasattr(self, "default") else value
+            value = self.attr_type.serialize(
+                value, **_get_context(self._attr_type_serialize_argspec, kwargs)
+            )
+            value = (
+                self._default() if value is None and hasattr(self, "default") else value
+            )
             if value in self.exclude:
                 raise ExcludedValueException()
             return value
@@ -282,7 +357,6 @@ def __init__(
                            the target resource.
         """
         if not types.Type.is_type(attr_type):
-
             if attr_type is not None:
                 attr = BYPASS
 
@@ -342,7 +416,9 @@ def __init__(self, attr_type=None, attr=None, required=True, curie=None):
         :param required: Is this list of links required to be present.
         :param curie: Link namespace prefix (e.g. "<prefix>:<name>") or Curie object.
         """
-        super(LinkList, self).__init__(attr_type=attr_type, attr=attr, required=required, curie=curie)
+        super(LinkList, self).__init__(
+            attr_type=attr_type, attr=attr, required=required, curie=curie
+        )
         self.attr_type = types.List(self.attr_type)
 
 
@@ -370,14 +446,22 @@ def __init__(self, name, href, templated=None, type=None):
 class Embedded(Attr):
     """Embedded attribute of schema."""
 
-    def __init__(self, attr_type: Union["halogen.Schema", "halogen.types.List"], attr=None, curie=None, required=True):
+    def __init__(
+        self,
+        attr_type: Union["halogen.Schema", "halogen.types.List"],
+        attr=None,
+        curie=None,
+        required=True,
+    ):
         """Embedded constructor.
 
         :param attr_type: Type, Schema or constant that does the type conversion of the attribute.
         :param attr: Attribute name, dot-separated attribute path or an `Accessor` instance.
         :param curie: The curie used for this embedded attribute.
         """
-        super(Embedded, self).__init__(attr_type=attr_type, attr=attr, required=required)
+        super(Embedded, self).__init__(
+            attr_type=attr_type, attr=attr, required=required
+        )
         self.curie = curie
         self.validate()
 
@@ -407,7 +491,9 @@ def validate(self):
         # Validate self link
         class_attributes = attribute_type.__dict__.get("__attrs__")
         if class_attributes is not None and "self" not in class_attributes.keys():
-            raise InvalidSchemaDefinition("Invalid HAL standard definition, need `self` link")
+            raise InvalidSchemaDefinition(
+                "Invalid HAL standard definition, need `self` link"
+            )
 
 
 class _Schema(types.Type):
@@ -469,7 +555,9 @@ def deserialize(cls, value, output=None, **kwargs):
                 errors.append(e)
             except (KeyError, AttributeError):
                 if attr.required:
-                    errors.append(exceptions.ValidationError("Missing attribute.", attr.name))
+                    errors.append(
+                        exceptions.ValidationError("Missing attribute.", attr.name)
+                    )
 
         if errors:
             raise exceptions.ValidationError(errors)
@@ -489,7 +577,9 @@ def __init__(cls, name, bases, clsattrs):
         cls.__class_attrs__ = OrderedDict()
         curies = set([])
 
-        attrs = [(key, value) for key, value in clsattrs.items() if isinstance(value, Attr)]
+        attrs = [
+            (key, value) for key, value in clsattrs.items() if isinstance(value, Attr)
+        ]
         attrs.sort(key=lambda attr: attr[1].creation_counter)
 
         # Collect the attributes and set their names.
@@ -508,7 +598,12 @@ def __init__(cls, name, bases, clsattrs):
 
         if curies:
             link = LinkList(
-                Schema(href=Attr(), name=Attr(), templated=Attr(required=False), type=Attr(required=False)),
+                Schema(
+                    href=Attr(),
+                    name=Attr(),
+                    templated=Attr(required=False),
+                    type=Attr(required=False),
+                ),
                 attr=lambda value: list(curies),
                 required=False,
             )
@@ -520,6 +615,108 @@ def __init__(cls, name, bases, clsattrs):
         for base in reversed(cls.__mro__):
             cls.__attrs__.update(getattr(base, "__class_attrs__", OrderedDict()))
 
+        cls.__output_type__ = _build_schema_output_type(cls)
+
+
+def _build_schema_output_type(schema_cls):
+    fields = {}
+    optional_fields = set()
+    compartment_fields = {}
+
+    for attr in schema_cls.__attrs__.values():
+        target = fields
+        if attr.compartment is not None:
+            target = compartment_fields.setdefault(attr.compartment, {})
+
+        key = attr.key
+        target[key] = _python_type_for_attr_type(attr.attr_type)
+        if not attr.required:
+            optional_fields.add((attr.compartment, key))
+
+    for compartment, comp_fields in compartment_fields.items():
+        fields[compartment] = _build_typed_dict_for_fields(
+            f"{schema_cls.__name__}{compartment.title().replace('_', '')}",
+            comp_fields,
+            {k for (c, k) in optional_fields if c == compartment},
+        )
+
+    return _build_typed_dict_for_fields(
+        f"{schema_cls.__name__}Serialized",
+        fields,
+        {k for (c, k) in optional_fields if c is None},
+    )
+
+
+def _build_typed_dict_for_fields(name, fields, optional_keys):
+    try:
+        from typing import NotRequired
+    except ImportError:
+        NotRequired = None
+
+    if NotRequired is not None and optional_keys:
+        annotations = {}
+        for key, value in fields.items():
+            if key in optional_keys:
+                annotations[key] = NotRequired[value]
+            else:
+                annotations[key] = value
+        return TypedDict(name, annotations, total=True)
+
+    total = not optional_keys
+    return TypedDict(name, dict(fields), total=total)
+
+
+def _python_type_for_attr_type(attr_type):
+    if isinstance(attr_type, halogen.schema._SchemaType):
+        return getattr(attr_type, "__output_type__", dict)
+
+    if isinstance(attr_type, halogen.types.List):
+        return TypingList[_python_type_for_attr_type(attr_type.item_type)]
+
+    if isinstance(attr_type, halogen.types.Nullable):
+        return Optional[_python_type_for_attr_type(attr_type.nested_type)]
+
+    if isinstance(attr_type, halogen.types.String):
+        return str
+
+    if isinstance(attr_type, halogen.types.Int):
+        return int
+
+    if isinstance(attr_type, halogen.types.Boolean):
+        return bool
+
+    if isinstance(
+        attr_type,
+        (
+            halogen.types.ISODateTime,
+            halogen.types.ISOUTCDateTime,
+            halogen.types.ISOUTCDate,
+        ),
+    ):
+        return str
+
+    if isinstance(attr_type, halogen.types.Amount):
+        return halogen.types.AmountSerialized
+
+    if isinstance(attr_type, halogen.types.Enum):
+        return str if not attr_type.use_values else object
+
+    if isinstance(attr_type, halogen.types.Type):
+        return Any
+
+    return type(attr_type)
+
+
+if TYPE_CHECKING:
+
+    class Schema(_Schema, metaclass=_SchemaType):
+        """Typing-only schema base class."""
+
+        @classmethod
+        def serialize(cls, value, **kwargs): ...
 
-Schema = _SchemaType("Schema", (_Schema,), {"__doc__": _Schema.__doc__})
-"""Schema is the basic class used for setting up schemas."""
+        @classmethod
+        def deserialize(cls, value, output=None, **kwargs): ...
+else:
+    Schema = _SchemaType("Schema", (_Schema,), {"__doc__": _Schema.__doc__})
+    """Schema is the basic class used for setting up schemas."""