Skip to content

Commit f6e36a9

Browse files
authored
Add eval_type_call which gets the result type of a function call at type check time. (#41)
For example, given: ```py def func[T](x: T) -> T: ... ``` we can evaluate the result type of when calling func on arguments of known types. ``` res = eval_type_call(func, int) assert res is int ```
1 parent 626d16f commit f6e36a9

4 files changed

Lines changed: 435 additions & 25 deletions

File tree

tests/test_qblike.py

Lines changed: 23 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,11 @@
22

33
from typing import Literal, Unpack
44

5-
from typemap.type_eval import eval_call, eval_typing
5+
from typemap.type_eval import (
6+
eval_call,
7+
eval_call_with_types,
8+
eval_typing,
9+
)
610
from typemap.typing import (
711
BaseTypedDict,
812
NewProtocol,
@@ -124,3 +128,21 @@ class select[...]:
124128
class PropsOnly[tests.test_qblike.Tgt]:
125129
name: tests.test_qblike.Property[str]
126130
""")
131+
132+
133+
def test_qblike_4():
134+
t = eval_call_with_types(
135+
select,
136+
A,
137+
x=bool,
138+
w=bool,
139+
z=bool,
140+
)
141+
fmt = format_helper.format_class(t)
142+
143+
assert fmt == textwrap.dedent("""\
144+
class select[...]:
145+
x: tests.test_qblike.Property[int]
146+
w: tests.test_qblike.Property[list[str]]
147+
z: tests.test_qblike.Link[tests.test_qblike.PropsOnly[tests.test_qblike.Tgt]]
148+
""")

tests/test_type_eval.py

Lines changed: 320 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,10 +17,11 @@
1717

1818
import pytest
1919

20-
from typemap.type_eval import eval_typing
20+
from typemap.type_eval import eval_call_with_types, eval_typing
2121
from typemap.typing import (
2222
Attrs,
2323
FromUnion,
24+
GenericCallable,
2425
GetArg,
2526
GetArgs,
2627
GetAttr,
@@ -1025,3 +1026,321 @@ def test_type_eval_annotated_03():
10251026
def test_type_eval_annotated_04():
10261027
res = eval_typing(GetAnnotations[GetAttr[AnnoTest, Literal["b"]]])
10271028
assert res == Literal["blah"]
1029+
1030+
1031+
def test_type_call_callable_01():
1032+
res = eval_call_with_types(Callable[[], int])
1033+
assert res is int
1034+
1035+
1036+
def test_type_call_callable_02():
1037+
res = eval_call_with_types(Callable[[Param[Literal["x"], int]], int], int)
1038+
assert res is int
1039+
1040+
1041+
def test_type_call_callable_03():
1042+
res = eval_call_with_types(
1043+
Callable[[Param[Literal["x"], int, Literal["keyword"]]], int], x=int
1044+
)
1045+
assert res is int
1046+
1047+
1048+
def test_type_call_callable_04():
1049+
class C: ...
1050+
1051+
res = eval_call_with_types(Callable[[Param[Literal["self"], Self]], int], C)
1052+
assert res is int
1053+
1054+
1055+
def test_type_call_callable_05():
1056+
class C: ...
1057+
1058+
res = eval_call_with_types(Callable[[Param[Literal["self"], Self]], C], C)
1059+
assert res is C
1060+
1061+
1062+
def test_type_call_callable_06():
1063+
class C: ...
1064+
1065+
res = eval_call_with_types(
1066+
Callable[[Param[Literal["self"], Self], Param[Literal["x"], int]], int],
1067+
C,
1068+
int,
1069+
)
1070+
assert res is int
1071+
1072+
1073+
def test_type_call_callable_07():
1074+
class C: ...
1075+
1076+
res = eval_call_with_types(
1077+
Callable[
1078+
[
1079+
Param[Literal["self"], Self],
1080+
Param[Literal["x"], int, Literal["keyword"]],
1081+
],
1082+
int,
1083+
],
1084+
C,
1085+
x=int,
1086+
)
1087+
assert res is int
1088+
1089+
1090+
def test_type_call_callable_08():
1091+
T = TypeVar("T")
1092+
res = eval_call_with_types(Callable[[Param[Literal["x"], T]], str], int)
1093+
assert res is str
1094+
1095+
1096+
def test_type_call_callable_09():
1097+
T = TypeVar("T")
1098+
res = eval_call_with_types(Callable[[Param[Literal["x"], T]], T], int)
1099+
assert res is int
1100+
1101+
1102+
def test_type_call_callable_10():
1103+
T = TypeVar("T")
1104+
1105+
class C(Generic[T]): ...
1106+
1107+
res = eval_call_with_types(Callable[[Param[Literal["x"], C[T]]], T], C[int])
1108+
assert res is int
1109+
1110+
1111+
def test_type_call_callable_11():
1112+
T = TypeVar("T")
1113+
1114+
class C(Generic[T]): ...
1115+
1116+
class D(C[int]): ...
1117+
1118+
class E(D): ...
1119+
1120+
res = eval_call_with_types(Callable[[Param[Literal["x"], C[T]]], T], D)
1121+
assert res is int
1122+
res = eval_call_with_types(Callable[[Param[Literal["x"], C[T]]], T], E)
1123+
assert res is int
1124+
1125+
1126+
def test_type_call_local_function_01():
1127+
def func(x: int) -> int: ...
1128+
1129+
res = eval_call_with_types(func, int)
1130+
assert res is int
1131+
1132+
1133+
def test_type_call_local_function_02():
1134+
def func(*, x: int) -> int: ...
1135+
1136+
res = eval_call_with_types(func, x=int)
1137+
assert res is int
1138+
1139+
1140+
def test_type_call_local_function_03():
1141+
def func[T](x: T) -> T: ...
1142+
1143+
res = eval_call_with_types(func, int)
1144+
assert res is int
1145+
1146+
1147+
def test_type_call_local_function_04():
1148+
class C: ...
1149+
1150+
def func(x: C) -> C: ...
1151+
1152+
res = eval_call_with_types(func, C)
1153+
assert res is C
1154+
1155+
1156+
def test_type_call_local_function_05():
1157+
class C: ...
1158+
1159+
def func[T](x: T) -> T: ...
1160+
1161+
res = eval_call_with_types(func, C)
1162+
assert res is C
1163+
1164+
1165+
def test_type_call_local_function_06():
1166+
T = TypeVar("T")
1167+
1168+
class C(Generic[T]): ...
1169+
1170+
def func[U](x: C[U]) -> C[U]: ...
1171+
1172+
res = eval_call_with_types(func, C[int])
1173+
assert res == C[int]
1174+
1175+
1176+
def test_type_call_local_function_07():
1177+
T = TypeVar("T")
1178+
1179+
class C(Generic[T]): ...
1180+
1181+
class D(C[int]): ...
1182+
1183+
class E(D): ...
1184+
1185+
def func[U](x: C[U]) -> U: ...
1186+
1187+
res = eval_call_with_types(func, D)
1188+
assert res is int
1189+
res = eval_call_with_types(func, E)
1190+
assert res is int
1191+
1192+
1193+
def test_type_call_local_function_08():
1194+
class C[T]: ...
1195+
1196+
class D(C[int]): ...
1197+
1198+
class E(C[str]): ...
1199+
1200+
class F(D, E): ...
1201+
1202+
def func[U](x: C[U]) -> U: ...
1203+
1204+
res = eval_call_with_types(func, F)
1205+
assert res is int
1206+
1207+
1208+
def test_type_call_local_function_09():
1209+
class C[T, U]: ...
1210+
1211+
def func[V](x: C[int, V]) -> V: ...
1212+
1213+
res = eval_call_with_types(func, C[int, str])
1214+
assert res is str
1215+
1216+
1217+
def test_type_call_bind_error_01():
1218+
T = TypeVar("T")
1219+
1220+
with pytest.raises(
1221+
ValueError, match="Type variable T is already bound to int, but got str"
1222+
):
1223+
eval_call_with_types(
1224+
Callable[[Param[Literal["x"], T], Param[Literal["y"], T]], T],
1225+
int,
1226+
str,
1227+
)
1228+
1229+
1230+
def test_type_call_bind_error_02():
1231+
def func[T](x: T, y: T) -> T: ...
1232+
1233+
with pytest.raises(
1234+
ValueError, match="Type variable T is already bound to int, but got str"
1235+
):
1236+
eval_call_with_types(func, int, str)
1237+
1238+
1239+
def test_type_call_bind_error_03():
1240+
T = TypeVar("T")
1241+
1242+
class C(Generic[T]): ...
1243+
1244+
with pytest.raises(
1245+
ValueError, match="Type variable T is already bound to int, but got str"
1246+
):
1247+
eval_call_with_types(
1248+
Callable[[Param[Literal["x"], C[T]], Param[Literal["y"], C[T]]], T],
1249+
C[int],
1250+
C[str],
1251+
)
1252+
1253+
1254+
def test_type_call_bind_error_04():
1255+
class C[T]: ...
1256+
1257+
def func[T](x: C[T], y: C[T]) -> T: ...
1258+
1259+
with pytest.raises(
1260+
ValueError, match="Type variable T is already bound to int, but got str"
1261+
):
1262+
eval_call_with_types(func, C[int], C[str])
1263+
1264+
1265+
def test_type_call_bind_error_05():
1266+
class C[T]: ...
1267+
1268+
class D[T]: ...
1269+
1270+
def func[T](x: C[T]) -> T: ...
1271+
1272+
with pytest.raises(ValueError, match="Argument type mismatch for x"):
1273+
eval_call_with_types(func, D[int])
1274+
1275+
1276+
type GetCallableMember[T, N: str] = GetArg[
1277+
tuple[
1278+
*[
1279+
GetType[m]
1280+
for m in Iter[Members[T]]
1281+
if (
1282+
IsSub[GetType[m], Callable]
1283+
or IsSub[GetType[m], GenericCallable]
1284+
)
1285+
and IsSub[GetName[m], N]
1286+
]
1287+
],
1288+
tuple,
1289+
0,
1290+
]
1291+
1292+
1293+
def test_type_call_member_01():
1294+
class C:
1295+
def invoke(self, x: int) -> int: ...
1296+
1297+
res = eval_call_with_types(GetCallableMember[C, Literal["invoke"]], C, int)
1298+
assert res is int
1299+
1300+
1301+
def test_type_call_member_02():
1302+
class C:
1303+
def invoke[T](self, x: T) -> T: ...
1304+
1305+
res = eval_call_with_types(GetCallableMember[C, Literal["invoke"]], C, int)
1306+
assert res is int
1307+
1308+
1309+
def test_type_call_member_03():
1310+
class C[T]:
1311+
def invoke(self, x: str) -> str: ...
1312+
1313+
res = eval_call_with_types(
1314+
GetCallableMember[C[int], Literal["invoke"]], C[int], str
1315+
)
1316+
assert res is str
1317+
1318+
1319+
def test_type_call_member_04():
1320+
class C[T]:
1321+
def invoke(self, x: T) -> T: ...
1322+
1323+
res = eval_call_with_types(
1324+
GetCallableMember[C[int], Literal["invoke"]], C[int], int
1325+
)
1326+
assert res is int
1327+
1328+
1329+
def test_type_call_member_05():
1330+
class C[T]:
1331+
def invoke(self) -> C[T]: ...
1332+
1333+
res = eval_call_with_types(
1334+
GetCallableMember[C[int], Literal["invoke"]], C[int]
1335+
)
1336+
assert res == C[int]
1337+
1338+
1339+
def test_type_call_member_06():
1340+
class C[T]:
1341+
def invoke[U](self, x: U) -> C[U]: ...
1342+
1343+
res = eval_call_with_types(
1344+
GetCallableMember[C[int], Literal["invoke"]], C[int], str
1345+
)
1346+
assert res == C[str]

typemap/type_eval/__init__.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@
77
from ._apply_generic import flatten_class
88

99
# XXX: this needs to go second due to nasty circularity -- try to fix that!!
10-
from ._eval_call import eval_call
10+
from ._eval_call import eval_call, eval_call_with_types
1111
from ._subtype import issubtype
1212
from ._subsim import issubsimilar
1313

@@ -19,6 +19,7 @@
1919
"eval_typing",
2020
"register_evaluator",
2121
"eval_call",
22+
"eval_call_with_types",
2223
"flatten_class",
2324
"issubtype",
2425
"issubsimilar",

0 commit comments

Comments
 (0)