diff --git a/tests/test_type_eval.py b/tests/test_type_eval.py index 15958f1..0963ab5 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,433 @@ 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(): + # 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: type[T], + ) -> AttrsAsSets[T]: + super().__init_subclass__() + + class B(A): + b: int + + 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_assignable_01(): + # Subclass retains information about bases + + class A: + a: int + + def __init_subclass__[T]( + cls: type[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: type[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'] diff --git a/typemap/type_eval/_eval_operators.py b/typemap/type_eval/_eval_operators.py index 4847b2a..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, ) @@ -43,6 +44,7 @@ SpecialFormEllipsis, StrConcat, Uncapitalize, + UpdateClass, Uppercase, _BoolLiteral, ) @@ -78,7 +80,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 +93,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 +127,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 +135,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 +164,100 @@ 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")) + ): + # 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 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 + ) + + # 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) + + # 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 + + +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 +595,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 +801,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 +809,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 +821,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 +851,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.