From 0b1cf8188e79440b022cf7d0c1f4700823d89ae5 Mon Sep 17 00:00:00 2001 From: dnwpark Date: Thu, 5 Feb 2026 15:51:15 -0800 Subject: [PATCH 1/4] Add tests. --- tests/test_type_eval.py | 360 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 360 insertions(+) diff --git a/tests/test_type_eval.py b/tests/test_type_eval.py index 15958f1..042ec70 100644 --- a/tests/test_type_eval.py +++ b/tests/test_type_eval.py @@ -47,6 +47,7 @@ Slice, SpecialFormEllipsis, StrConcat, + UpdateClass, Uppercase, ) @@ -1655,6 +1656,365 @@ def static_method(x: int) -> int: ... """) +def test_update_class_members_01(): + # Non-generic UpdateClass + class A: + a1: int # not overridden + a2: int # overridden + + def __init_subclass__( + cls, + ) -> UpdateClass[ + Member[Literal["a2"], str], + Member[Literal["b1"], str], + Member[Literal["b2"], str], + ]: + super().__init_subclass__() + + def f(self) -> int: ... + + class B(A): + b0: int # omitted + b1: int # overridden + # b2 added in UpdateClass + + def g(self) -> int: ... # omitted + + # Attrs + attrs = eval_typing(Attrs[B]) + assert ( + attrs + == tuple[ + Member[Literal["a1"], int, Never, Never, A], + Member[Literal["a2"], str, Never, Never, B], + Member[Literal["b1"], str, Never, Never, B], + Member[Literal["b2"], str, Never, Never, B], + ] + ) + + # Members + members = eval_typing(Members[B]) + assert ( + members + == tuple[ + Member[Literal["a1"], int, Never, Never, A], + Member[Literal["a2"], str, Never, Never, B], + Member[Literal["b1"], str, Never, Never, B], + Member[Literal["b2"], str, Never, Never, B], + Member[ + Literal["__init_subclass__"], + classmethod[ + A, + tuple[()], + UpdateClass[ + Member[Literal["a2"], str], + Member[Literal["b1"], str], + Member[Literal["b2"], str], + ], + ], + Literal["ClassVar"], + object, + A, + ], + Member[ + Literal["f"], + Callable[[Param[Literal["self"], A]], int], + Literal["ClassVar"], + object, + A, + ], + ] + ) + + # GetMember + assert members == eval_typing( + tuple[ + GetMember[B, Literal["a1"]], + GetMember[B, Literal["a2"]], + GetMember[B, Literal["b1"]], + GetMember[B, Literal["b2"]], + GetMember[B, Literal["__init_subclass__"]], + GetMember[B, Literal["f"]], + ] + ) + m = eval_typing(GetMember[B, Literal["g"]]) + assert m == Never + + +type MembersExceptInitSubclass[T] = tuple[ + *[ + m + for m in Iter[Members[T]] + if not IsAssignable[GetName[m], Literal["__init_subclass__"]] + ] +] + + +def test_update_class_members_02(): + # Generic UpdateClass, does not use T + class A: + a1: int # not overridden + a2: int # overridden + + def __init_subclass__[T]( + cls: type[T], + ) -> UpdateClass[ + Member[Literal["a2"], str], + Member[Literal["b1"], str], + Member[Literal["b2"], str], + ]: + super().__init_subclass__() + + def f(self) -> int: ... + + class B(A): + b0: int # omitted + b1: int # overridden + # b2 added in UpdateClass + + def g(self) -> int: ... # omitted + + # Attrs + attrs = eval_typing(Attrs[B]) + assert ( + attrs + == tuple[ + Member[Literal["a1"], int, Never, Never, A], + Member[Literal["a2"], str, Never, Never, B], + Member[Literal["b1"], str, Never, Never, B], + Member[Literal["b2"], str, Never, Never, B], + ] + ) + + # Members + members = eval_typing(MembersExceptInitSubclass[B]) + assert ( + members + == tuple[ + Member[Literal["a1"], int, Never, Never, A], + Member[Literal["a2"], str, Never, Never, B], + Member[Literal["b1"], str, Never, Never, B], + Member[Literal["b2"], str, Never, Never, B], + Member[ + Literal["f"], + Callable[[Param[Literal["self"], A]], int], + Literal["ClassVar"], + object, + A, + ], + ] + ) + + # GetMember + assert members == eval_typing( + tuple[ + GetMember[B, Literal["a1"]], + GetMember[B, Literal["a2"]], + GetMember[B, Literal["b1"]], + GetMember[B, Literal["b2"]], + GetMember[B, Literal["f"]], + ] + ) + m = eval_typing(GetMember[B, Literal["g"]]) + assert m == Never + + +def test_update_class_members_03(): + # Generic UpdateClass, uses T + type AttrsAsSets[T] = UpdateClass[ + *[Member[GetName[m], set[GetType[m]]] for m in Iter[Attrs[T]]] + ] + + class A: + a: int + + def __init_subclass__[T]( + cls: type[T], + ) -> AttrsAsSets[T]: + super().__init_subclass__() + + def f(self) -> int: ... + + class B(A): + b: str + + def g(self) -> int: ... # omitted + + # Attrs + attrs = eval_typing(Attrs[B]) + assert ( + attrs + == tuple[ + Member[Literal["a"], set[int], Never, Never, B], + Member[Literal["b"], set[str], Never, Never, B], + ] + ) + + # Members + members = eval_typing(MembersExceptInitSubclass[B]) + assert ( + members + == tuple[ + Member[Literal["a"], set[int], Never, Never, B], + Member[Literal["b"], set[str], Never, Never, B], + Member[ + Literal["f"], + Callable[[Param[Literal["self"], A]], int], + Literal["ClassVar"], + object, + A, + ], + ] + ) + + # GetMember + assert members == eval_typing( + tuple[ + GetMember[B, Literal["a"]], + GetMember[B, Literal["b"]], + GetMember[B, Literal["f"]], + ] + ) + m = eval_typing(GetMember[B, Literal["g"]]) + assert m == Never + + +def test_update_class_members_04(): + # Non-UpdateClass __init_subclass__ unaffected + class A: + a: int + + def __init_subclass__( + cls, + ) -> None: + super().__init_subclass__() + + def f(self) -> int: ... + + class B(A): + b: int + + def g(self) -> int: ... + + # Attrs + attrs = eval_typing(Attrs[B]) + assert ( + attrs + == tuple[ + Member[Literal["a"], int, Never, Never, A], + Member[Literal["b"], int, Never, Never, B], + ] + ) + + # Members + members = eval_typing(Members[B]) + assert ( + members + == tuple[ + Member[Literal["a"], int, Never, Never, A], + Member[Literal["b"], int, Never, Never, B], + Member[ + Literal["__init_subclass__"], + classmethod[ + A, + tuple[()], + None, + ], + Literal["ClassVar"], + object, + A, + ], + Member[ + Literal["f"], + Callable[[Param[Literal["self"], A]], int], + Literal["ClassVar"], + object, + A, + ], + Member[ + Literal["g"], + Callable[[Param[Literal["self"], B]], int], + Literal["ClassVar"], + object, + B, + ], + ] + ) + + # GetMember + assert members == eval_typing( + tuple[ + GetMember[B, Literal["a"]], + GetMember[B, Literal["b"]], + GetMember[B, Literal["__init_subclass__"]], + GetMember[B, Literal["f"]], + GetMember[B, Literal["g"]], + ] + ) + + +def test_update_class_inheritance_01(): + # __init_subclass__ calls follow normal MRO + class A: + a: int + + def __init_subclass__[T]( + cls: T, + ) -> UpdateClass[()]: ... + + class B(A): + b: int + + class C(B): + c: int + + +def test_update_class_is_sub_01(): + # Subclass retains information about bases + + class A: + a: int + + def __init_subclass__[T]( + cls: T, + ) -> UpdateClass[()]: + super().__init_subclass__() + + class B(A): + b: int + + class C(B): + c: int + + assert eval_typing(IsAssignable[B, A]) == _BoolLiteral[True] + assert eval_typing(IsAssignable[C, B]) == _BoolLiteral[True] + assert eval_typing(IsAssignable[C, A]) == _BoolLiteral[True] + + +def test_update_class_getarg_01(): + # Subclass is able to get args from base classes + + class A[X0, X1]: + a: X1 + + def __init_subclass__[T]( + cls: T, + ) -> UpdateClass[()]: + super().__init_subclass__() + + class B[Y](A[int, Y]): + b: Y + + class C(B[float]): + c: float + + assert eval_typing(GetArg[B[str], A, Literal[0]]) is int + assert eval_typing(GetArg[B[str], A, Literal[1]]) is str + assert eval_typing(GetArg[C, B, Literal[0]]) is float + assert eval_typing(GetArg[C, A, Literal[0]]) is int + assert eval_typing(GetArg[C, A, Literal[1]]) is float + + ############## type XTest[X] = Annotated[X, 'blah'] From 29a999cafdc8e3d7ef50b65a63d0a1869014cc06 Mon Sep 17 00:00:00 2001 From: dnwpark Date: Thu, 5 Feb 2026 15:52:11 -0800 Subject: [PATCH 2/4] First implementation. --- typemap/type_eval/_eval_operators.py | 117 ++++++++++++++++++++++++--- typemap/typing.py | 4 + 2 files changed, 108 insertions(+), 13 deletions(-) diff --git a/typemap/type_eval/_eval_operators.py b/typemap/type_eval/_eval_operators.py index 4847b2a..62887c1 100644 --- a/typemap/type_eval/_eval_operators.py +++ b/typemap/type_eval/_eval_operators.py @@ -43,6 +43,7 @@ SpecialFormEllipsis, StrConcat, Uncapitalize, + UpdateClass, Uppercase, _BoolLiteral, ) @@ -78,7 +79,7 @@ def _make_init_type(v): return typing.Literal[(v,)] -def get_annotated_type_hints(cls, **kwargs): +def get_annotated_type_hints(cls, *, ctx, **kwargs): """Get the type hints/quals for a cls annotated with definition site. This traverses the mro and finds the definition site for each annotation. @@ -91,6 +92,11 @@ def get_annotated_type_hints(cls, **kwargs): for abox in reversed(box.mro): acls = abox.alias_type() + if abox is box and (updated_cls := _eval_init_subclass(box, ctx)): + # For the class itself, apply all UpdateClass from + # ancesstors' __init_subclass__ to get the final type. + abox = _apply_generic.box(updated_cls) + annos, _ = _apply_generic.get_local_defns(abox) for k, ty in annos.items(): quals = set() @@ -120,7 +126,7 @@ def get_annotated_type_hints(cls, **kwargs): return hints -def get_annotated_method_hints(cls): +def get_annotated_method_hints(cls, *, ctx): # TODO: Cache the box (slash don't need it??) box = _apply_generic.box(cls) @@ -128,6 +134,11 @@ def get_annotated_method_hints(cls): for abox in reversed(box.mro): acls = abox.alias_type() + if abox is box and (updated_cls := _eval_init_subclass(box, ctx)): + # For the class itself, apply all UpdateClass from + # ancesstors' __init_subclass__ to get the final type. + abox = _apply_generic.box(updated_cls) + _, dct = _apply_generic.get_local_defns(abox) for name, attr in dct.items(): if isinstance( @@ -152,6 +163,81 @@ def get_annotated_method_hints(cls): return hints +def _eval_init_subclass( + box: _apply_generic.Boxed, ctx: typing.Any +) -> type | None: + """Get type after all __init_subclass__ with UpdateClass are evaluated.""" + for abox in reversed(box.mro[1:]): # Skip the type itself + if ms := _get_update_class_members(box.cls, abox.alias_type(), ctx=ctx): + return _create_updated_class(box.cls, ms, ctx=ctx) + + return None + + +def _get_update_class_members( + cls: type, base: type, ctx: typing.Any +) -> list[Member] | None: + if ( + (init_subclass := base.__dict__.get("__init_subclass__")) + and ( + init_subclass_annos := getattr( + init_subclass, "__annotations__", None + ) + ) + and (ret_annotation := init_subclass_annos.get("return")) + and _typing_inspect.is_generic_alias(ret_annotation) + and typing.get_origin(ret_annotation) is UpdateClass + ): + # Substitute the cls type var with the current class + # This may not happen if cls is not generic! + if ( + cls_anno := next( + (v for k, v in init_subclass_annos.items() if k != "return"), + None, + ) + ) and isinstance(cls_anno, typing.TypeVar): + substitution = {cls_anno: cls} + ret_annotation = _apply_generic.substitute( + ret_annotation, substitution + ) + + evaled_ret = _eval_types(ret_annotation, ctx=ctx) + + return [m for m in typing.get_args(evaled_ret)] + + return None + + +def _create_updated_class(t: type, ms: list[Member], ctx) -> type: + dct: dict[str, object] = {} + + # Copy the module + dct["__module__"] = t.__module__ + + # Process the new members from UpdateClass + dct["__annotations__"] = annos = {} + for m in ms: + tname, typ, quals, init, _ = typing.get_args(m) + member_name = _eval_literal(tname, ctx) + typ = _eval_types(typ, ctx) + tquals = _eval_types(quals, ctx) + + if type_eval.issubtype( + typing.Literal["ClassVar"], tquals + ) and _is_method_like(typ): + dct[member_name] = _callable_type_to_method(member_name, typ, ctx) + else: + # Update/add the annotation + annos[member_name] = _add_quals(typ, tquals) + _unpack_init(dct, member_name, init) + + # Create the updated class + mcls = type(t) + cls = mcls(t.__name__, t.__bases__, dct) + + return cls + + def _union_elems(tp, ctx): tp = _eval_types(tp, ctx) if tp is typing.Never: @@ -489,10 +575,15 @@ def _callable_type_to_method(name, typ, ctx): has_pos_only = any(_is_pos_only(p) for p in typing.get_args(params)) quals = typing.Literal["positional"] if has_pos_only else typing.Never # Override the receiver type with type[Self]. - # An annoying thing to know is that for a member classmethod of C, - # cls *should* be type[C], but if it was not explicitly annotated, it - # will be C. - cls_param = Param[typing.Literal["cls"], type[typing.Self], quals] + if name == "__init_subclass__" and isinstance(cls, typing.TypeVar): + # For __init_subclass__ generic on cls: T, keep type[T] + cls_typ = type[cls] # type: ignore[name-defined] + else: + # An annoying thing to know is that for a member classmethod of C, + # cls *should* be type[C], but if it was not explicitly annotated, + # it will be C. + cls_typ = type[typing.Self] # type: ignore[name-defined] + cls_param = Param[typing.Literal["cls"], cls_typ, quals] typ = typing.Callable[[cls_param] + list(typing.get_args(params)), ret] elif head is staticmethod: params, ret = typing.get_args(typ) @@ -690,7 +781,7 @@ def _hints_to_members(hints, ctx): @type_eval.register_evaluator(Attrs) @_lift_over_unions def _eval_Attrs(tp, *, ctx): - hints = get_annotated_type_hints(tp, include_extras=True) + hints = get_annotated_type_hints(tp, include_extras=True, ctx=ctx) return _hints_to_members(hints, ctx) @@ -698,8 +789,8 @@ def _eval_Attrs(tp, *, ctx): @_lift_over_unions def _eval_Members(tp, *, ctx): hints = { - **get_annotated_type_hints(tp, include_extras=True), - **get_annotated_method_hints(tp), + **get_annotated_type_hints(tp, include_extras=True, ctx=ctx), + **get_annotated_method_hints(tp, ctx=ctx), } return _hints_to_members(hints, ctx) @@ -710,8 +801,8 @@ def _eval_GetMember(tp, prop, *, ctx): # XXX: extras? name = _from_literal(prop) hints = { - **get_annotated_type_hints(tp, include_extras=True), - **get_annotated_method_hints(tp), + **get_annotated_type_hints(tp, include_extras=True, ctx=ctx), + **get_annotated_method_hints(tp, ctx=ctx), } if name in hints: return _hint_to_member(name, *hints[name], ctx=ctx) @@ -740,8 +831,8 @@ def _eval_GetMemberType(tp, prop, *, ctx): # XXX: extras? name = _from_literal(prop) hints = { - **get_annotated_type_hints(tp, include_extras=True), - **get_annotated_method_hints(tp), + **get_annotated_type_hints(tp, include_extras=True, ctx=ctx), + **get_annotated_method_hints(tp, ctx=ctx), } if name in hints: return hints[name][0] diff --git a/typemap/typing.py b/typemap/typing.py index 6fc94f3..aac61c7 100644 --- a/typemap/typing.py +++ b/typemap/typing.py @@ -247,6 +247,10 @@ class NewProtocol[*T]: pass +class UpdateClass[*Ms]: + pass + + class RaiseError[S: str, *Ts]: """Raise a type error with the given message when evaluated. From 87ddb5a8182224f208f337ce4dec70f775cd133f Mon Sep 17 00:00:00 2001 From: dnwpark Date: Thu, 5 Feb 2026 18:25:08 -0800 Subject: [PATCH 3/4] Fix generic UpdateClass. --- typemap/type_eval/_eval_operators.py | 38 +++++++++++++++++++++------- 1 file changed, 29 insertions(+), 9 deletions(-) diff --git a/typemap/type_eval/_eval_operators.py b/typemap/type_eval/_eval_operators.py index 62887c1..38a5587 100644 --- a/typemap/type_eval/_eval_operators.py +++ b/typemap/type_eval/_eval_operators.py @@ -12,6 +12,7 @@ from typemap import type_eval from typemap.type_eval import _apply_generic, _typing_inspect from typemap.type_eval._eval_typing import ( + _child_context, _eval_types, _get_class_type_hint_namespaces, ) @@ -185,25 +186,44 @@ def _get_update_class_members( ) ) and (ret_annotation := init_subclass_annos.get("return")) - and _typing_inspect.is_generic_alias(ret_annotation) - and typing.get_origin(ret_annotation) is UpdateClass ): # Substitute the cls type var with the current class # This may not happen if cls is not generic! if ( - cls_anno := next( - (v for k, v in init_subclass_annos.items() if k != "return"), - None, + ( + cls_anno := next( + ( + v + for k, v in init_subclass_annos.items() + if k != "return" + ), + None, + ) ) - ) and isinstance(cls_anno, typing.TypeVar): - substitution = {cls_anno: cls} + and typing.get_origin(cls_anno) is type + and (cls_type_args := typing.get_args(cls_anno)) + and (cls_type := cls_type_args[0]) + and isinstance(cls_type, typing.TypeVar) + ): + substitution = {cls_type: cls} ret_annotation = _apply_generic.substitute( ret_annotation, substitution ) - evaled_ret = _eval_types(ret_annotation, ctx=ctx) + # Evaluate the return annotation + # Do it in a child context, so the evaluations are isolated. For + # example, if the return annotation uses Attrs[MyClass], we want + # Attrs[MyClass] to be evaluated with the updated class, not the + # original. + with _child_context() as ctx: + evaled_ret = _eval_types(ret_annotation, ctx=ctx) - return [m for m in typing.get_args(evaled_ret)] + # If the result is an UpdateClass, return the members + if ( + _typing_inspect.is_generic_alias(evaled_ret) + and typing.get_origin(evaled_ret) is UpdateClass + ): + return [m for m in typing.get_args(evaled_ret)] return None From 8c3eb5b61cefae2e3d2ec0b726f471196593acd3 Mon Sep 17 00:00:00 2001 From: dnwpark Date: Thu, 5 Feb 2026 18:40:46 -0800 Subject: [PATCH 4/4] Fix tests. --- tests/test_type_eval.py | 80 +++++++++++++++++++++++++++++++++++++---- 1 file changed, 74 insertions(+), 6 deletions(-) diff --git a/tests/test_type_eval.py b/tests/test_type_eval.py index 042ec70..0963ab5 100644 --- a/tests/test_type_eval.py +++ b/tests/test_type_eval.py @@ -1954,29 +1954,97 @@ def g(self) -> int: ... def test_update_class_inheritance_01(): + # current class init subclass is not applied + type AttrsAsSets[T] = UpdateClass[ + *[Member[GetName[m], set[GetType[m]]] for m in Iter[Attrs[T]]] + ] + + class A: + a: int + + def __init_subclass__[T]( + cls: type[T], + ) -> AttrsAsSets[T]: + super().__init_subclass__() + + class B(A): + b: int + + def __init_subclass__[T]( + cls: type[T], + ) -> AttrsAsSets[T]: + super().__init_subclass__() + + attrs = eval_typing(Attrs[B]) + assert ( + attrs + == tuple[ + Member[Literal["a"], set[int], Never, Never, B], + Member[Literal["b"], set[int], Never, Never, B], + ] + ) + + +def test_update_class_inheritance_02(): # __init_subclass__ calls follow normal MRO + type AttrsAsSets[T] = UpdateClass[ + *[Member[GetName[m], set[GetType[m]]] for m in Iter[Attrs[T]]] + ] + type AttrsAsList[T] = UpdateClass[ + *[Member[GetName[m], list[GetType[m]]] for m in Iter[Attrs[T]]] + ] + type AttrsAsTuple[T] = UpdateClass[ + *[Member[GetName[m], tuple[GetType[m]]] for m in Iter[Attrs[T]]] + ] + class A: a: int def __init_subclass__[T]( - cls: T, - ) -> UpdateClass[()]: ... + cls: type[T], + ) -> AttrsAsSets[T]: + super().__init_subclass__() class B(A): b: int - class C(B): + def __init_subclass__[T]( + cls: type[T], + ) -> AttrsAsList[T]: + super().__init_subclass__() + + class C: c: int + def __init_subclass__[T]( + cls: type[T], + ) -> AttrsAsTuple[T]: + super().__init_subclass__() + + class D(B, C): + d: int + + attrs = eval_typing(Attrs[D]) + # MRO = D, B, A, C, object + assert ( + attrs + == tuple[ + Member[Literal["c"], tuple[set[list[int]]], Never, Never, D], + Member[Literal["a"], tuple[set[list[int]]], Never, Never, D], + Member[Literal["b"], tuple[set[list[int]]], Never, Never, D], + Member[Literal["d"], tuple[set[list[int]]], Never, Never, D], + ] + ) + -def test_update_class_is_sub_01(): +def test_update_class_is_assignable_01(): # Subclass retains information about bases class A: a: int def __init_subclass__[T]( - cls: T, + cls: type[T], ) -> UpdateClass[()]: super().__init_subclass__() @@ -1998,7 +2066,7 @@ class A[X0, X1]: a: X1 def __init_subclass__[T]( - cls: T, + cls: type[T], ) -> UpdateClass[()]: super().__init_subclass__()