Support Annotated and write bigger examples

spec-draft.rst

@@ -191,6 +191,28 @@ We can put more in, but this is what typescript has.
 We can actually implement the case functions in terms of them and a
 bunch of conditionals, but shouldn't (especially if we want it to work for all unicode!).
 
+---------
+Annotated
+---------
+
+Libraries like FastAPI use annotations heavily, and we would like to be able to use annotations to drive type-level computation decision making.
+
+We understand that this may be controversial, as currently Annotated may be fully ignored by typecheckers. The operations proposed are:
+
+* ``GetAnnotations[T]`` - Fetch the annotations of a potentially Annotated type, as Literals. Examples::
+
+    GetAnnotations[Annotated[int, 'xxx']] = Literal['xxx']
+    GetAnnotations[Annotated[int, 'xxx', 5]] = Literal['xxx', 5]
+    GetAnnotations[int] = Never
+
+
+* ``DropAnnotations[T]`` - Drop the annotations of a potentially Annotated type. Examples::
+
+    DropAnnotations[Annotated[int, 'xxx']] = int
+    DropAnnotations[Annotated[int, 'xxx', 5]] = int
+    DropAnnotations[int] = int
+
+
 -------------------
 String manipulation
 -------------------

tests/test_fastapilike_1.py

@@ -0,0 +1,265 @@
+import dataclasses
+import enum
+import textwrap
+
+from typing import Annotated, Callable, Literal, Union, Self
+
+from typemap.type_eval import eval_typing
+from typemap.typing import (
+    NewProtocol,
+    Iter,
+    Attrs,
+    Is,
+    GetAnnotations,
+    DropAnnotations,
+    FromUnion,
+    GetType,
+    GetName,
+    GetQuals,
+    Member,
+    Members,
+    Param,
+)
+
+from . import format_helper
+
+
+class PropQuals(enum.StrEnum):
+    HIDDEN = "HIDDEN"
+    PRIMARY = "PRIMARY"
+    HAS_DEFAULT = "HAS_DEFAULT"
+
+
+@dataclasses.dataclass(frozen=True)
+class _Default:
+    val: object
+
+
+type Hidden[T] = Annotated[T, Literal[PropQuals.HIDDEN]]
+type Primary[T] = Annotated[T, Literal[PropQuals.PRIMARY]]
+type HasDefault[T, default] = Annotated[
+    T, _Default(default), Literal[PropQuals.HAS_DEFAULT]
+]
+
+
+####
+
+type InitFnType[T] = Member[
+    Literal["__init__"],
+    Callable[
+        [
+            Param[Literal["self"], Self],
+            *[
+                Param[
+                    GetName[p],
+                    DropAnnotations[GetType[p]],
+                    Literal["keyword", "default"]
+                    if Is[
+                        Literal[PropQuals.HAS_DEFAULT],
+                        GetAnnotations[GetType[p]],
+                    ]
+                    else Literal["keyword"],
+                ]
+                for p in Iter[Attrs[T]]
+            ],
+        ],
+        None,
+    ],
+    Literal["ClassVar"],
+]
+type AddInit[T] = NewProtocol[
+    InitFnType[T],
+    *[x for x in Iter[Members[T]]],
+]
+
+"""TODO:
+
+We would really like to instead write:
+
+type AddInit[T] = NewProtocol[
+    InitFnType[T],
+    *Members[T]],
+]
+
+but we struggle here because typing wants to unpack the Members tuple
+itself.  I'm not sure if there is a nice way to resolve this. We
+*could* make our consumers (NewProtocol etc) be more flexible about
+these things but I don't think that is right.
+
+The frustrating thing is that it doesn't do much with the unpacked
+version, just some checks!
+
+We could fix typing to allow it, and probably provide a hack around it
+in the mean time.
+
+Lurr! Writing *this* gets past the typing checks (though we don't
+support it yet):
+
+type AddInit[T] = NewProtocol[
+    InitFnType[T],
+    *tuple[*Members[T]],
+]
+"""
+
+type AllOptional[T] = NewProtocol[
+    *[
+        Member[GetName[p], GetType[p] | None, GetQuals[p]]
+        for p in Iter[Attrs[T]]
+    ]
+]
+
+# Strip `| None` from a type by iterating over its union components
+# and filtering
+type NotOptional[T] = Union[*[x for x in Iter[FromUnion[T]] if not Is[x, None]]]
+
+# Adjust an attribute type for use in Public below by dropping | None for
+# primary keys and stripping all annotations.
+type FixPublicType[T] = DropAnnotations[
+    NotOptional[T] if Is[Literal[PropQuals.PRIMARY], GetAnnotations[T]] else T
+]
+
+# Strip out everything that is Hidden and also make the primary key required
+# Drop all the annotations, since this is for data getting returned to users
+# from the DB, so we don't need default values.
+type Public[T] = NewProtocol[
+    *[
+        Member[GetName[p], FixPublicType[GetType[p]], GetQuals[p]]
+        for p in Iter[Attrs[T]]
+        if not Is[Literal[PropQuals.HIDDEN], GetAnnotations[GetType[p]]]
+    ]
+]
+
+# Create takes everything but the primary key and preserves defaults
+type Create[T] = NewProtocol[
+    *[
+        Member[GetName[p], GetType[p], GetQuals[p]]
+        for p in Iter[Attrs[T]]
+        if not Is[Literal[PropQuals.PRIMARY], GetAnnotations[GetType[p]]]
+    ]
+]
+
+
+# Update takes everything but the primary key, but makes them all have
+# None defaults
+type Update[T] = NewProtocol[
+    *[
+        Member[
+            GetName[p],
+            HasDefault[DropAnnotations[GetType[p]] | None, None],
+            GetQuals[p],
+        ]
+        for p in Iter[Attrs[T]]
+        if not Is[Literal[PropQuals.PRIMARY], GetAnnotations[GetType[p]]]
+    ]
+]
+
+
+####
+
+# This is the FastAPI example code that we are trying to repair!
+# Adapted from https://fastapi.tiangolo.com/tutorial/sql-databases/#heroupdate-the-data-model-to-update-a-hero
+"""
+class HeroBase(SQLModel):
+    name: str = Field(index=True)
+    age: int | None = Field(default=None, index=True)
+
+
+class Hero(HeroBase, table=True):
+    id: int | None = Field(default=None, primary_key=True)
+    secret_name: str
+
+
+class HeroPublic(HeroBase):
+    id: int
+
+
+class HeroCreate(HeroBase):
+    secret_name: str
+
+
+class HeroUpdate(HeroBase):
+    name: str | None = None
+    age: int | None = None
+    secret_name: str | None = None
+"""
+
+
+class Hero:
+    id: Primary[
+        HasDefault[int | None, None]
+    ]  # = Field(default=None, primary_key=True)
+
+    name: str
+    age: HasDefault[int | None, None]  # = Field(default=None, index=True)
+
+    secret_name: Hidden[str]
+
+
+#######
+
+
+def test_eval_drop_optional_1():
+    tgt = eval_typing(NotOptional[int | None])
+    assert tgt is int
+
+
+def test_fastapi_like_1():
+    tgt = eval_typing(Public[Hero])
+    fmt = format_helper.format_class(tgt)
+
+    assert fmt == textwrap.dedent("""\
+        class Public[tests.test_fastapilike_1.Hero]:
+            id: int
+            name: str
+            age: int | None
+    """)
+
+
+def test_fastapi_like_2():
+    tgt = eval_typing(Create[Hero])
+    fmt = format_helper.format_class(tgt)
+
+    assert fmt == textwrap.dedent("""\
+        class Create[tests.test_fastapilike_1.Hero]:
+            name: str
+            age: typing.Annotated[int | None, _Default(val=None), typing.Literal[<PropQuals.HAS_DEFAULT: 'HAS_DEFAULT'>]]
+            secret_name: typing.Annotated[str, typing.Literal[<PropQuals.HIDDEN: 'HIDDEN'>]]
+    """)
+
+
+def test_fastapi_like_3():
+    tgt = eval_typing(Update[Hero])
+    fmt = format_helper.format_class(tgt)
+
+    assert fmt == textwrap.dedent("""\
+        class Update[tests.test_fastapilike_1.Hero]:
+            name: typing.Annotated[str | None, _Default(val=None), typing.Literal[<PropQuals.HAS_DEFAULT: 'HAS_DEFAULT'>]]
+            age: typing.Annotated[int | None, _Default(val=None), typing.Literal[<PropQuals.HAS_DEFAULT: 'HAS_DEFAULT'>]]
+            secret_name: typing.Annotated[str | None, _Default(val=None), typing.Literal[<PropQuals.HAS_DEFAULT: 'HAS_DEFAULT'>]]
+    """)
+
+
+def test_fastapi_like_4():
+    tgt = eval_typing(AddInit[Public[Hero]])
+    fmt = format_helper.format_class(tgt)
+
+    assert fmt == textwrap.dedent("""\
+        class AddInit[tests.test_fastapilike_1.Public[tests.test_fastapilike_1.Hero]]:
+            id: int
+            name: str
+            age: int | None
+            def __init__(self: Self, *, id: int, name: str, age: int | None) -> None: ...
+    """)
+
+
+def test_fastapi_like_6():
+    tgt = eval_typing(AddInit[Update[Hero]])
+    fmt = format_helper.format_class(tgt)
+
+    assert fmt == textwrap.dedent("""\
+        class AddInit[tests.test_fastapilike_1.Update[tests.test_fastapilike_1.Hero]]:
+            name: typing.Annotated[str | None, _Default(val=None), typing.Literal[<PropQuals.HAS_DEFAULT: 'HAS_DEFAULT'>]]
+            age: typing.Annotated[int | None, _Default(val=None), typing.Literal[<PropQuals.HAS_DEFAULT: 'HAS_DEFAULT'>]]
+            secret_name: typing.Annotated[str | None, _Default(val=None), typing.Literal[<PropQuals.HAS_DEFAULT: 'HAS_DEFAULT'>]]
+            def __init__(self: Self, *, name: str | None = ..., age: int | None = ..., secret_name: str | None = ...) -> None: ...
+    """)

tests/test_type_eval.py

@@ -2,6 +2,7 @@
 import textwrap
 import unittest
 from typing import (
+    Annotated,
     Any,
     Callable,
     Generic,
@@ -24,6 +25,7 @@
     GetAttr,
     GetName,
     GetType,
+    GetAnnotations,
     Is,
     Iter,
     Length,
@@ -807,3 +809,38 @@ def test_callable_to_signature():
         '(_arg0: int, /, b: int, c: int = ..., *args: int, '
         'd: int, e: int = ..., **kwargs: int) -> int'
     )
+
+
+##############
+
+type XTest[X] = Annotated[X, 'blah']
+
+
+class AnnoTest:
+    a: XTest[int]
+    b: XTest[Literal["test"]]
+
+
+def test_type_eval_annotated_01():
+    res = format_helper.format_class(eval_typing(AnnoTest))
+
+    assert res == textwrap.dedent("""\
+        class AnnoTest:
+            a: typing.Annotated[int, 'blah']
+            b: typing.Annotated[typing.Literal['test'], 'blah']
+    """)
+
+
+def test_type_eval_annotated_02():
+    res = eval_typing(Is[GetAttr[AnnoTest, Literal["a"]], int])
+    assert res is True
+
+
+def test_type_eval_annotated_03():
+    res = eval_typing(Uppercase[GetAttr[AnnoTest, Literal["b"]]])
+    assert res == Literal["TEST"]
+
+
+def test_type_eval_annotated_04():
+    res = eval_typing(GetAnnotations[GetAttr[AnnoTest, Literal["b"]]])
+    assert res == Literal["blah"]

typemap/type_eval/_apply_generic.py

@@ -195,6 +195,9 @@ def _get_closure_types(af: types.FunctionType) -> dict[str, type]:
     }
 
 
+EXCLUDED_ATTRIBUTES = typing.EXCLUDED_ATTRIBUTES - {'__init__'}  # type: ignore[attr-defined]
+
+
 def get_local_defns(boxed: Boxed) -> tuple[dict[str, Any], dict[str, Any]]:
     annos: dict[str, Any] = {}
     dct: dict[str, Any] = {}
@@ -238,7 +241,7 @@ def get_local_defns(boxed: Boxed) -> tuple[dict[str, Any], dict[str, Any]]:
         annos.update(af)
 
     for name, orig in boxed.cls.__dict__.items():
-        if name in typing.EXCLUDED_ATTRIBUTES:  # type: ignore[attr-defined]
+        if name in EXCLUDED_ATTRIBUTES:
             continue
 
         stuff = inspect.unwrap(orig)