Add an additional argument to GetArg to represent where to look

Author: msullivan

spec-draft.rst

@@ -18,9 +18,7 @@ Big Q: what should be an error and what should return Never?
         # *[... for t in ...] arguments
         | <ident>[<variadic-type-arg(<type>)> +]
 
-        # This is syntax because taking an int literal makes it a
-        # special form.
-        | GetArg[<type>, <int-literal>]
+        | <string-or-int-literal>  # Only accepted in arguments to new functions?
 
 
    # Type conditional checks are just boolean compositions of
@@ -48,12 +46,15 @@ Big Q: what should be an error and what should return Never?
          if <type-bool>
 
 
-``type-for(T)`` is a parameterized grammar rule, which can take
-different types. Not sure if we actually need this though---now it is
-only used for Any/All.
+``type-for(T)`` is a parameterized grammar rule, which can take different types. Not sure if we actually need this though---now it is only used for Any/All.
 
 ---
 
+* ``GetArg[T, Base, Idx: Literal[str]]`` - returns the type argument number ``Idx`` to ``T`` when interpreted as ``Base``, or ``Never`` if it cannot be. (That is, if we have  ``class A(B[int]): ...``, then ``GetArg[A, B, 0] == int``.)
+* ``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.
+
+
 # TODO: NewProtocol needs a way of doing bases also...
 # TODO: New TypedDict setup
 
@@ -77,10 +78,6 @@ only used for Any/All.
 
 # TODO: how to deal with special forms like Callable and tuple[T, ...]
 
-* ``GetArgs[T]`` - returns a tuple containing all of the type arguments
-* ``FromUnion[T]`` - returns a tuple containing all of the union
-  elements, or a 1-ary tuple containing T if it is not a union.
-
 # TODO: How to do IsUnion? Might need a ``Length`` for tuples?
 
 

tests/test_qblike.py

@@ -31,7 +31,7 @@ class Link[T]:
 ]
 
 # Conditional type alias!
-type FilterLinks[T] = Link[PropsOnly[GetArg[T, 0]]] if Is[T, Link] else T
+type FilterLinks[T] = Link[PropsOnly[GetArg[T, Link, 0]]] if Is[T, Link] else T
 
 
 # Basic filtering

tests/test_type_dir.py

@@ -1,10 +1,11 @@
 import textwrap
-from typing import TypeVar, Literal, Union
+from typing import Never, Literal, Union, TypeVar
 
 from typemap.type_eval import eval_typing
 from typemap.typing import (
     NewProtocol,
     Member,
+    GetArg,
     GetName,
     GetType,
     Iter,
@@ -65,10 +66,14 @@ class Last[O]:
     last: O | Literal[True]
 
 
+# XXX: the mixing of int and float as arguments to Base is questionable!
 class Final(Mine, Ordinary, Wrapper[float], AnotherBase[float], Last[int]):
     pass
 
 
+type BaseArg[T] = GetArg[T, Base, 0] if Is[T, Base] else Never
+
+
 type AllOptional[T] = NewProtocol[
     *[Member[GetName[p], GetType[p] | None] for p in Iter[Attrs[T]]]
 ]
@@ -239,3 +244,8 @@ def test_type_dir_7():
 typemap.typing.Param[typing.Literal['b'], int, typing.Literal['=']]], \
 dict[str, int]], typing.Literal['ClassVar'], typing.Never]"
     )
+
+
+def test_type_dir_8():
+    d = eval_typing(BaseArg[Final])
+    assert d is int

typemap/type_eval/_apply_generic.py

@@ -226,4 +226,5 @@ def apply(cls: type[Any]) -> dict[str, Any]:
         dct[k] = _eval_typing.eval_typing(v)
 
     dct["__annotations__"] = annos
+    dct["__generalized_mro__"] = mro_boxed
     return dct

typemap/type_eval/_typing_inspect.py

@@ -125,6 +125,15 @@ def is_literal(t: Any) -> bool:
     return is_generic_alias(t) and get_origin(t) is Literal  # type: ignore [comparison-overlap]
 
 
+def get_head(t: Any) -> type | None:
+    if is_generic_alias(t):
+        return get_origin(t)
+    elif isinstance(t, type):
+        return t
+    else:
+        return None
+
+
 def is_eval_proxy(t: Any) -> TypeGuard[type[_eval_typing._EvalProxy]]:
     return isinstance(t, type) and issubclass(t, _eval_typing._EvalProxy)
 

typemap/typing.py

@@ -80,10 +80,10 @@ class Member[N: str, T, Q: str = typing.Never, D = typing.Never]:
     pass
 
 
-type GetName[T: Member] = GetArg[T, 0]  # type: ignore[valid-type]
-type GetType[T: Member] = GetArg[T, 1]  # type: ignore[valid-type]
-type GetQuals[T: Member] = GetArg[T, 2]  # type: ignore[valid-type]
-type GetDefiner[T: Member] = GetArg[T, 3]  # type: ignore[valid-type]
+type GetName[T: Member] = GetArg[T, Member, 0]  # type: ignore[valid-type]
+type GetType[T: Member] = GetArg[T, Member, 1]  # type: ignore[valid-type]
+type GetQuals[T: Member] = GetArg[T, Member, 2]  # type: ignore[valid-type]
+type GetDefiner[T: Member] = GetArg[T, Member, 3]  # type: ignore[valid-type]
 
 
 ##################################################################
@@ -233,12 +233,42 @@ def GetAttr(self, arg):
     return typing.get_type_hints(type_eval.eval_typing(lhs))[name]
 
 
+def _get_args(tp, base) -> typing.Any:
+    # XXX: check against base!!
+    evaled = type_eval.eval_typing(tp)
+
+    tp_head = _typing_inspect.get_head(tp)
+    base_head = _typing_inspect.get_head(base)
+    # XXX: not sure this is what we want!
+    # at the very least we want unions I think
+    if not tp_head or not base_head:
+        return None
+
+    if tp_head is base_head:
+        return typing.get_args(evaled)
+
+    # Scan the fully-annotated MRO to find the base
+    elif gen_mro := getattr(evaled, "__generalized_mro__", None):
+        for box in gen_mro:
+            if box.cls is base_head:
+                return tuple(box.args.values())
+        return None
+
+    else:
+        # or error??
+        return None
+
+
 @_SpecialForm
-def GetArg(self, arg):
-    tp, idx = arg
-    args = typing.get_args(type_eval.eval_typing(tp))
+def GetArg(self, arg) -> typing.Any:
+    # XXX: Unions
+    tp, base, idx = arg
+    args = _get_args(tp, base)
+    if args is None:
+        return typing.Never
+
     try:
-        return args[idx]
+        return args[_from_literal(idx)]
     except IndexError:
         return typing.Never