Implement GetArgs

Author: msullivan

spec-draft.rst

@@ -69,7 +69,7 @@ It's important that there be a clearly specified type language for the type-leve
   Special forms unfortunately require some special handling: the arguments list of a ``Callable`` will be packed in a tuple, and a ``...`` will become ``SpecialFormEllipsis``.
 
 
-* ``GetArgs[T, Base]`` - returns a tuple containing all of the type arguments of ``T`` when interpreted as ``Base``, or ``Never`` if it cannot be. (TODO: UNIMPLEMENTED)
+* ``GetArgs[T, Base]`` - returns a tuple containing all of the type arguments of ``T`` when interpreted as ``Base``, or ``Never`` if it cannot be.
 * ``FromUnion[T]`` - returns a tuple containing all of the union elements, or a 1-ary tuple containing T if it is not a union.
 
 

tests/test_type_eval.py

@@ -7,6 +7,7 @@
     Attrs,
     FromUnion,
     GetArg,
+    GetArgs,
     GetAttr,
     GetName,
     GetType,
@@ -178,6 +179,16 @@ def test_getarg_never():
     assert d is Never
 
 
+def test_eval_getargs():
+    t = dict[int, str]
+    args = eval_typing(GetArgs[t, dict])
+    assert args == tuple[int, str]
+
+    t = dict
+    args = eval_typing(GetArgs[t, dict])
+    assert args == tuple[Any, Any]
+
+
 def test_eval_getarg_callable():
     # oh hmmmmmmm -- yeah maybe callable could be fully bespoke if we
     # disallowed putting Callable here...!

typemap/type_eval/_eval_operators.py

@@ -16,6 +16,7 @@
     Capitalize,
     FromUnion,
     GetArg,
+    GetArgs,
     GetAttr,
     IsSubSimilar,
     IsSubtype,
@@ -247,12 +248,11 @@ def _eval_GetAttr(lhs, prop, *, ctx):
     return typing.get_type_hints(lhs)[name]
 
 
-def _get_args(tp, base, ctx) -> typing.Any:
+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)
-    base_head = _typing_inspect.get_head(base)
     if not tp_head or not base_head:
         return None
 
@@ -271,6 +271,14 @@ def _get_args(tp, base, ctx) -> typing.Any:
         return None
 
 
+def _get_args(tp, base, ctx) -> typing.Any:
+    base_head = _typing_inspect.get_head(base)
+    args = _get_raw_args(tp, base, ctx)
+    if args == ():
+        args = _get_defaults(base_head)
+    return args
+
+
 def _fix_type(tp):
     """Fix up a type getting returned from GetArg
 
@@ -383,8 +391,6 @@ def _get_defaults(base_head):
 def _eval_GetArg(tp, base, idx, *, ctx) -> typing.Any:
     base_head = _typing_inspect.get_head(base)
     args = _get_args(tp, base_head, ctx)
-    if args == ():
-        args = _get_defaults(base_head)
     if args is None:
         return typing.Never
 
@@ -394,6 +400,16 @@ def _eval_GetArg(tp, base, idx, *, ctx) -> typing.Any:
         return typing.Never
 
 
+@type_eval.register_evaluator(GetArgs)
+@_lift_over_unions
+def _eval_GetArgs(tp, base, *, ctx) -> typing.Any:
+    base_head = _typing_inspect.get_head(base)
+    args = _get_args(tp, base_head, ctx)
+    if args is None:
+        return typing.Never
+    return tuple[*args]  # type: ignore[valid-type]
+
+
 @type_eval.register_evaluator(Length)
 @_lift_over_unions
 def _eval_Length(tp, *, ctx) -> typing.Any:

typemap/typing.py

@@ -52,6 +52,10 @@ class GetArg[Tp, Base, Idx: int]:
     pass
 
 
+class GetArgs[Tp, Base]:
+    pass
+
+
 class Length[S: tuple]:
     pass