Skip to content

Commit dff3f20

Browse files
authored
Ensure that generic self works for methods and class methods. (#97)
Mostly already worked, just needed to pass the receiver for class methods to ensure things like this work: ```py @classmethod def f[T]( cls: type[T], ) -> ... ```
1 parent 46ab27d commit dff3f20

File tree

2 files changed

+208
-4
lines changed

2 files changed

+208
-4
lines changed

tests/test_type_eval.py

Lines changed: 202 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@
2121

2222
import pytest
2323

24-
from typemap.type_eval import eval_typing
24+
from typemap.type_eval import _ensure_context, eval_typing
2525
from typemap.typing import _BoolLiteral
2626

2727
from typemap_extensions import (
@@ -508,6 +508,207 @@ def f(self, x): ...
508508
)
509509

510510

511+
def test_getmember_05():
512+
# member method, generic self, complex return type
513+
class A:
514+
def member_method[T](self: T) -> GetMemberType[T, Literal["x"]]: ...
515+
516+
class B(A):
517+
x: int
518+
519+
class C(A):
520+
x: str
521+
522+
m = eval_typing(GetMember[A, Literal["member_method"]])
523+
assert eval_typing(GetName[m]) == Literal["member_method"]
524+
assert eval_typing(IsAssignable[GetType[m], GenericCallable])
525+
assert eval_typing(GetQuals[m]) == Literal["ClassVar"]
526+
assert eval_typing(GetDefiner[m]) == A
527+
528+
ft = m.__args__[1].__args__[1]
529+
with _ensure_context():
530+
assert (
531+
eval_typing(ft(A)) == Callable[[Param[Literal["self"], A]], Never]
532+
)
533+
assert eval_typing(ft(B)) == Callable[[Param[Literal["self"], B]], int]
534+
assert eval_typing(ft(C)) == Callable[[Param[Literal["self"], C]], str]
535+
536+
537+
def test_getmember_06():
538+
# member method, generic self, complex param type
539+
class A:
540+
def member_method[T](
541+
self: T, x: GetMemberType[T, Literal["x"]]
542+
) -> None: ...
543+
544+
class B(A):
545+
x: int
546+
547+
class C(A):
548+
x: str
549+
550+
m = eval_typing(GetMember[A, Literal["member_method"]])
551+
assert eval_typing(GetName[m]) == Literal["member_method"]
552+
assert eval_typing(IsAssignable[GetType[m], GenericCallable])
553+
assert eval_typing(GetQuals[m]) == Literal["ClassVar"]
554+
assert eval_typing(GetDefiner[m]) == A
555+
556+
ft = m.__args__[1].__args__[1]
557+
with _ensure_context():
558+
assert (
559+
eval_typing(ft(A))
560+
== Callable[
561+
[Param[Literal["self"], A], Param[Literal["x"], Never]], None
562+
]
563+
)
564+
assert (
565+
eval_typing(ft(B))
566+
== Callable[
567+
[Param[Literal["self"], B], Param[Literal["x"], int]], None
568+
]
569+
)
570+
assert (
571+
eval_typing(ft(C))
572+
== Callable[
573+
[Param[Literal["self"], C], Param[Literal["x"], str]], None
574+
]
575+
)
576+
577+
578+
def test_getmember_07():
579+
# class method, generic self, complex return type
580+
class A:
581+
@classmethod
582+
def class_method[T](cls: type[T]) -> GetMemberType[T, Literal["x"]]: ...
583+
584+
class B(A):
585+
x: int
586+
587+
class C(A):
588+
x: str
589+
590+
m = eval_typing(GetMember[A, Literal["class_method"]])
591+
assert eval_typing(GetName[m]) == Literal["class_method"]
592+
assert eval_typing(IsAssignable[GetType[m], GenericCallable])
593+
assert eval_typing(GetQuals[m]) == Literal["ClassVar"]
594+
assert eval_typing(GetDefiner[m]) == A
595+
596+
ft = m.__args__[1].__args__[1]
597+
with _ensure_context():
598+
assert eval_typing(ft(A)) == classmethod[A, tuple[()], Never]
599+
assert eval_typing(ft(B)) == classmethod[B, tuple[()], int]
600+
assert eval_typing(ft(C)) == classmethod[C, tuple[()], str]
601+
602+
603+
def test_getmember_08():
604+
# class method, generic self, complex param type
605+
class A:
606+
@classmethod
607+
def class_method[T](
608+
cls: type[T], x: GetMemberType[T, Literal["x"]]
609+
) -> None: ...
610+
611+
class B(A):
612+
x: int
613+
614+
class C(A):
615+
x: str
616+
617+
m = eval_typing(GetMember[A, Literal["class_method"]])
618+
assert eval_typing(GetName[m]) == Literal["class_method"]
619+
assert eval_typing(IsAssignable[GetType[m], GenericCallable])
620+
assert eval_typing(GetQuals[m]) == Literal["ClassVar"]
621+
assert eval_typing(GetDefiner[m]) == A
622+
623+
ft = m.__args__[1].__args__[1]
624+
with _ensure_context():
625+
assert (
626+
eval_typing(ft(A))
627+
== classmethod[A, tuple[Param[Literal["x"], Never]], None]
628+
)
629+
assert (
630+
eval_typing(ft(B))
631+
== classmethod[B, tuple[Param[Literal["x"], int]], None]
632+
)
633+
assert (
634+
eval_typing(ft(C))
635+
== classmethod[C, tuple[Param[Literal["x"], str]], None]
636+
)
637+
638+
639+
def test_getmember_09():
640+
# member method, generic self, iterating over members
641+
class A:
642+
def f[T](
643+
self: T,
644+
) -> tuple[
645+
*[
646+
GetType[m]
647+
for m in Iter[Attrs[T]]
648+
if not IsAssignable[
649+
Slice[GetName[m], None, Literal[1]], Literal["_"]
650+
]
651+
]
652+
]: ...
653+
654+
class B(A):
655+
a: bool
656+
b: str
657+
_c: int
658+
_d: float
659+
660+
m = eval_typing(GetMember[A, Literal["f"]])
661+
assert eval_typing(GetName[m]) == Literal["f"]
662+
assert eval_typing(IsAssignable[GetType[m], GenericCallable])
663+
assert eval_typing(GetQuals[m]) == Literal["ClassVar"]
664+
assert eval_typing(GetDefiner[m]) == A
665+
666+
ft = m.__args__[1].__args__[1]
667+
with _ensure_context():
668+
assert (
669+
eval_typing(ft(A))
670+
== Callable[[Param[Literal["self"], A]], tuple[()]]
671+
)
672+
assert (
673+
eval_typing(ft(B))
674+
== Callable[[Param[Literal["self"], B]], tuple[bool, str]]
675+
)
676+
677+
678+
def test_getmember_10():
679+
# class method, generic self, iterating over members
680+
class A:
681+
@classmethod
682+
def f[T](
683+
cls: type[T],
684+
) -> tuple[
685+
*[
686+
GetType[m]
687+
for m in Iter[Attrs[T]]
688+
if not IsAssignable[
689+
Slice[GetName[m], None, Literal[1]], Literal["_"]
690+
]
691+
]
692+
]: ...
693+
694+
class B(A):
695+
a: bool
696+
b: str
697+
_x: int
698+
_y: float
699+
700+
m = eval_typing(GetMember[B, Literal["f"]])
701+
assert eval_typing(GetName[m]) == Literal["f"]
702+
assert eval_typing(IsAssignable[GetType[m], GenericCallable])
703+
assert eval_typing(GetQuals[m]) == Literal["ClassVar"]
704+
assert eval_typing(GetDefiner[m]) == A
705+
706+
ft = m.__args__[1].__args__[1]
707+
with _ensure_context():
708+
assert eval_typing(ft(A)) == classmethod[A, tuple[()], tuple[()]]
709+
assert eval_typing(ft(B)) == classmethod[B, tuple[()], tuple[bool, str]]
710+
711+
511712
def test_getarg_never():
512713
d = eval_typing(GetArg[Never, object, Literal[0]])
513714
assert d is Never

typemap/type_eval/_apply_generic.py

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -347,8 +347,9 @@ def get_local_defns(
347347
if stuck and stuff.__type_params__:
348348
type_params = stuff.__type_params__
349349
str_args = boxed.str_args
350+
canonical_cls = boxed.canonical_cls
350351

351-
def _make_lambda(fn, o, sa, tp):
352+
def _make_lambda(fn, o, sa, tp, cls):
352353
from ._eval_operators import _function_type_from_sig
353354

354355
def lam(*vs):
@@ -362,14 +363,16 @@ def lam(*vs):
362363
)
363364
sig = _resolved_function_signature(fn, args)
364365
return _function_type_from_sig(
365-
sig, o, receiver_type=None
366+
sig, o, receiver_type=cls
366367
)
367368

368369
return lam
369370

370371
gc = GenericCallable[ # type: ignore[valid-type,misc]
371372
tuple[*type_params], # type: ignore[valid-type]
372-
_make_lambda(stuff, orig, str_args, type_params),
373+
_make_lambda(
374+
stuff, orig, str_args, type_params, canonical_cls
375+
),
373376
]
374377
annos[name] = typing.ClassVar[gc]
375378
elif local_fn is not None:

0 commit comments

Comments
 (0)