Properly handle Quals in NewProtocol and Members

Author: msullivan

tests/test_type_dir.py

@@ -1,24 +1,25 @@
 import textwrap
-from typing import Never, Literal, Union, TypeVar
+import typing
+from typing import Literal, Never, TypeVar, Union
 
 from typemap.type_eval import eval_typing
 from typemap.typing import (
-    NewProtocol,
-    Member,
+    Attrs,
+    FromUnion,
     GetArg,
     GetName,
+    GetQuals,
     GetType,
+    Is,
     Iter,
-    Attrs,
+    Member,
     Members,
-    FromUnion,
+    NewProtocol,
     Uppercase,
-    Is,
 )
 
 from . import format_helper
 
-
 type OrGotcha[K] = K | Literal["gotcha!"]
 
 type StrForInt[X] = (str | OrGotcha[X]) if X is int else (X | OrGotcha[X])
@@ -39,6 +40,8 @@ class Base[T]:
     t: dict[str, StrForInt[T]]
     kkk: K
 
+    fin: typing.Final[int]
+
     def foo(self, a: T | None, b: int = 0) -> dict[str, T]:
         pass
 
@@ -75,14 +78,20 @@ class Final(Mine, Ordinary, Wrapper[float], AnotherBase[float], Last[int]):
 
 
 type AllOptional[T] = NewProtocol[
-    *[Member[GetName[p], GetType[p] | None] for p in Iter[Attrs[T]]]
+    *[
+        Member[GetName[p], GetType[p] | None, GetQuals[p]]
+        for p in Iter[Attrs[T]]
+    ]
 ]
 
 type OptionalFinal = AllOptional[Final]
 
 
 type Capitalize[T] = NewProtocol[
-    *[Member[Uppercase[GetName[p]], GetType[p]] for p in Iter[Attrs[T]]]
+    *[
+        Member[Uppercase[GetName[p]], GetType[p], GetQuals[p]]
+        for p in Iter[Attrs[T]]
+    ]
 ]
 
 type Prims[T] = NewProtocol[
@@ -102,6 +111,7 @@ class Final(Mine, Ordinary, Wrapper[float], AnotherBase[float], Last[int]):
                     if not Is[t, Literal]
                 ]
             ],
+            GetQuals[p],
         ]
         for p in Iter[Attrs[T]]
     ]
@@ -137,6 +147,7 @@ class Final(Mine, Ordinary, Wrapper[float], AnotherBase[float], Last[int]):
                     if not Is[IsLiteral[t], Literal[True]]
                 ]
             ],
+            GetQuals[p],
         ]
         for p in Iter[Attrs[T]]
     ]
@@ -152,6 +163,7 @@ class Final:
             iii: str | int | typing.Literal['gotcha!']
             t: dict[str, str | int | typing.Literal['gotcha!']]
             kkk: ~K
+            fin: typing.Final[int]
             x: tests.test_type_dir.Wrapper[int | None]
             ordinary: str
             def foo(self, a: int | None, b: int = 0) -> dict[str, int]: ...
@@ -172,6 +184,7 @@ class AllOptional[tests.test_type_dir.Final]:
             iii: str | int | typing.Literal['gotcha!'] | None
             t: dict[str, str | int | typing.Literal['gotcha!']] | None
             kkk: ~K | None
+            fin: typing.Final[int | None]
             x: tests.test_type_dir.Wrapper[int | None] | None
             ordinary: str | None
     """)
@@ -186,6 +199,7 @@ class Capitalize[tests.test_type_dir.Final]:
             III: str | int | typing.Literal['gotcha!']
             T: dict[str, str | int | typing.Literal['gotcha!']]
             KKK: ~K
+            FIN: typing.Final[int]
             X: tests.test_type_dir.Wrapper[int | None]
             ORDINARY: str
     """)
@@ -197,6 +211,7 @@ def test_type_dir_4():
     assert format_helper.format_class(d) == textwrap.dedent("""\
         class Prims[tests.test_type_dir.Final]:
             last: int | typing.Literal[True]
+            fin: typing.Final[int]
             ordinary: str
     """)
 
@@ -211,6 +226,7 @@ class NoLiterals1[tests.test_type_dir.Final]:
             iii: str | int
             t: dict[str, str | int | typing.Literal['gotcha!']]
             kkk: ~K
+            fin: typing.Final[int]
             x: tests.test_type_dir.Wrapper[int | None]
             ordinary: str
     """)
@@ -225,6 +241,7 @@ class NoLiterals2[tests.test_type_dir.Final]:
             iii: str | int
             t: dict[str, str | int | typing.Literal['gotcha!']]
             kkk: ~K
+            fin: typing.Final[int]
             x: tests.test_type_dir.Wrapper[int | None]
             ordinary: str
     """)

typemap/type_eval/_eval_operators.py

@@ -56,8 +56,24 @@ def get_annotated_type_hints(cls, **kwargs):
             continue
         for k in acls.__annotations__:
             if k not in hints:
-                # XXX: TODO: Strip ClassVar/Final
-                hints[k] = ohints[k], (), acls
+                quals = set()
+                ty = ohints[k]
+
+                # Strip ClassVar/Final from ty and add them to quals
+                while True:
+                    for form in [typing.ClassVar, typing.Final]:
+                        if _typing_inspect.is_special_form(ty, form):
+                            quals.add(form.__name__)
+                            ty = (
+                                typing.get_args(ty)[0]
+                                if typing.get_args(ty)
+                                else typing.Any
+                            )
+                            break
+                    else:
+                        break
+
+                hints[k] = ty, tuple(sorted(quals)), acls
 
         # Stop early if we are done.
         if len(hints) == len(ohints):
@@ -249,7 +265,6 @@ def _eval_GetAttr(lhs, prop, *, ctx):
 
 
 def _get_raw_args(tp, base_head, ctx) -> typing.Any:
-    # XXX: check against base!!
     evaled = _eval_types(tp, ctx)
 
     tp_head = _typing_inspect.get_head(tp)
@@ -444,15 +459,23 @@ def func(*args, ctx):
 ##################################################################
 
 
+def _add_quals(typ, quals):
+    for qual in (typing.ClassVar, typing.Final):
+        if type_eval.issubsimilar(typing.Literal[qual.__name__], quals):
+            typ = qual[typ]
+    return typ
+
+
 @type_eval.register_evaluator(NewProtocol)
 def _eval_NewProtocol(*etyps: Member, ctx):
     dct: dict[str, object] = {}
     dct["__annotations__"] = {
         # XXX: Should eval_typing on the etyps evaluate the arguments??
-        _from_literal(typing.get_args(prop)[0], ctx): _eval_types(
-            typing.get_args(prop)[1], ctx
+        _from_literal(name, ctx): _add_quals(
+            _eval_types(typ, ctx),
+            _eval_types(quals, ctx),
         )
-        for prop in etyps
+        for name, typ, quals, _ in (typing.get_args(prop) for prop in etyps)
     }
 
     module_name = __name__

typemap/type_eval/_typing_inspect.py

@@ -4,28 +4,37 @@
 
 
 import typing
-
-from typing import (
+from types import GenericAlias, UnionType
+from typing import (  # type: ignore [attr-defined]  # noqa: PLC2701
     Annotated,
     Any,
-    ClassVar,
     ForwardRef,
     Literal,
     TypeGuard,
     TypeVar,
     Union,
+    _GenericAlias,
+    _SpecialGenericAlias,
     get_args,
     get_origin,
 )
-from typing import _GenericAlias, _SpecialGenericAlias  # type: ignore [attr-defined]  # noqa: PLC2701
+
 from typing_extensions import TypeAliasType, TypeVarTuple, Unpack
-from types import GenericAlias, UnionType
 
 from . import _eval_typing
 
 
-def is_classvar(t: Any) -> bool:
-    return t is ClassVar or (is_generic_alias(t) and get_origin(t) is ClassVar)  # type: ignore [comparison-overlap]
+def is_special_form(t: Any, form: Any) -> bool:
+    """Check if t is a special form or a generic alias of that form.
+
+    Args:
+        t: The type to check
+        form: The special form to check against (e.g., ClassVar, Final, Literal)
+
+    Returns:
+        True if t is the special form or a generic alias with that origin
+    """
+    return t is form or (is_generic_alias(t) and get_origin(t) is form)  # type: ignore [comparison-overlap]
 
 
 def is_generic_alias(t: Any) -> TypeGuard[GenericAlias]:
@@ -142,12 +151,12 @@ def is_eval_proxy(t: Any) -> TypeGuard[type[_eval_typing._EvalProxy]]:
 
 __all__ = (
     "is_annotated",
-    "is_classvar",
     "is_forward_ref",
     "is_generic_alias",
     "is_generic_type_alias",
     "is_literal",
     "is_optional_type",
+    "is_special_form",
     "is_type_alias",
     "is_union_type",
 )