Handle default type params

Author: msullivan

spec-draft.rst

@@ -65,8 +65,11 @@ It's important that there be a clearly specified type language for the type-leve
 
 ---
 
-* ``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[C]): ...``, then ``GetArg[A, B, 0] == C`` while ``GetArg[A, A, 0] == Never``)
-* ``GetArgs[T, Base]`` - returns a tuple containing all of the type arguments of ``T`` when interpreted as ``Base``, or ``Never`` if it cannot be.
+* ``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[C]): ...``, then ``GetArg[A, B, 0] == C`` while ``GetArg[A, A, 0] == Never``).
+  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)
 * ``FromUnion[T]`` - returns a tuple containing all of the union elements, or a 1-ary tuple containing T if it is not a union.
 
 
@@ -92,8 +95,6 @@ It's important that there be a clearly specified type language for the type-leve
 * ``GetAttr[T, S: Literal[str]]``
   TODO: How should GetAttr interact with descriptors/classmethod? I am leaning towards it should apply the descriptor...
 
-# TODO: how to deal with special forms like Callable and tuple[T, ...]
-
 * ``Length[T: tuple]`` - get the length of a tuple as an int literal (...or ``Literal[None]`` if it is unbounded)
 
 String manipulation operations for string Literal types.

tests/test_type_eval.py

@@ -1,6 +1,6 @@
 import textwrap
 import unittest
-from typing import Literal, Never, Tuple
+from typing import Any, Callable, Generic, List, Literal, Never, Tuple, TypeVar
 
 from typemap.type_eval import eval_typing
 from typemap.typing import (
@@ -15,6 +15,7 @@
     Length,
     Member,
     NewProtocol,
+    SpecialFormEllipsis,
     StrConcat,
     StrSlice,
     Uppercase,
@@ -177,6 +178,211 @@ def test_getarg_never():
     assert d is Never
 
 
+def test_eval_getarg_callable():
+    # oh hmmmmmmm -- yeah maybe callable could be fully bespoke if we
+    # disallowed putting Callable here...!
+    t = Callable[[int, str], str]
+    args = eval_typing(GetArg[t, Callable, 0])
+    assert args == tuple[int, str]
+
+    t = Callable[int, str]
+    args = eval_typing(GetArg[t, Callable, 0])
+    assert args == tuple[int]
+
+    t = Callable[[], str]
+    args = eval_typing(GetArg[t, Callable, 0])
+    assert args == tuple[()]
+
+    t = Callable[..., str]
+    args = eval_typing(GetArg[t, Callable, 0])
+    assert args == SpecialFormEllipsis
+
+    t = Callable
+    args = eval_typing(GetArg[t, Callable, 0])
+    assert args == SpecialFormEllipsis
+
+    t = Callable
+    args = eval_typing(GetArg[t, Callable, 1])
+    assert args == Any
+
+
+def test_eval_getarg_tuple():
+    t = tuple[int, ...]
+    args = eval_typing(GetArg[t, tuple, 1])
+    assert args == SpecialFormEllipsis
+
+    t = tuple
+    args = eval_typing(GetArg[t, tuple, 0])
+    assert args == Any
+
+    args = eval_typing(GetArg[t, tuple, 1])
+    assert args == SpecialFormEllipsis
+
+
+def test_eval_getarg_list():
+    t = list[int]
+    arg = eval_typing(GetArg[t, list, 0])
+    assert arg is int
+
+    t = List[int]
+    arg = eval_typing(GetArg[t, list, 0])
+    assert arg is int
+
+    t = list
+    arg = eval_typing(GetArg[t, list, 0])
+    assert arg == Any
+
+    t = List
+    arg = eval_typing(GetArg[t, list, 0])
+    assert arg == Any
+
+    t = list[int]
+    arg = eval_typing(GetArg[t, List, 0])
+    assert arg is int
+
+    t = List[int]
+    arg = eval_typing(GetArg[t, List, 0])
+    assert arg is int
+
+    t = list
+    arg = eval_typing(GetArg[t, List, 0])
+    assert arg == Any
+
+    t = List
+    arg = eval_typing(GetArg[t, List, 0])
+    assert arg == Any
+
+    # indexing with -1 equivalent to 0
+    t = list[int]
+    arg = eval_typing(GetArg[t, list, -1])
+    assert arg is int
+
+    t = List[int]
+    arg = eval_typing(GetArg[t, list, -1])
+    assert arg is int
+
+    t = list
+    arg = eval_typing(GetArg[t, list, -1])
+    assert arg == Any
+
+    t = List
+    arg = eval_typing(GetArg[t, list, -1])
+    assert arg == Any
+
+    t = list[int]
+    arg = eval_typing(GetArg[t, List, -1])
+    assert arg is int
+
+    t = List[int]
+    arg = eval_typing(GetArg[t, List, -1])
+    assert arg is int
+
+    t = list
+    arg = eval_typing(GetArg[t, List, -1])
+    assert arg == Any
+
+    t = List
+    arg = eval_typing(GetArg[t, List, -1])
+    assert arg == Any
+
+    # indexing with 1 always fails
+    t = list[int]
+    arg = eval_typing(GetArg[t, list, 1])
+    assert arg == Never
+
+    t = List[int]
+    arg = eval_typing(GetArg[t, list, 1])
+    assert arg == Never
+
+    t = list
+    arg = eval_typing(GetArg[t, list, 1])
+    assert arg == Never
+
+    t = List
+    arg = eval_typing(GetArg[t, list, 1])
+    assert arg == Never
+
+    t = list[int]
+    arg = eval_typing(GetArg[t, List, 1])
+    assert arg == Never
+
+    t = List[int]
+    arg = eval_typing(GetArg[t, List, 1])
+    assert arg == Never
+
+    t = list
+    arg = eval_typing(GetArg[t, List, 1])
+    assert arg == Never
+
+    t = List
+    arg = eval_typing(GetArg[t, List, 1])
+    assert arg == Never
+
+
+def test_eval_getarg_custom_01():
+    class A[T]:
+        pass
+
+    t = A[int]
+    assert eval_typing(GetArg[t, A, 0]) is int
+    assert eval_typing(GetArg[t, A, -1]) is int
+    assert eval_typing(GetArg[t, A, 1]) == Never
+
+    t = A
+    assert eval_typing(GetArg[t, A, 0]) == Any
+    assert eval_typing(GetArg[t, A, -1]) == Any
+    assert eval_typing(GetArg[t, A, 1]) == Never
+
+
+def test_eval_getarg_custom_02():
+    T = TypeVar("T")
+
+    class A(Generic[T]):
+        pass
+
+    t = A[int]
+    assert eval_typing(GetArg[t, A, 0]) is int
+    assert eval_typing(GetArg[t, A, -1]) is int
+    assert eval_typing(GetArg[t, A, 1]) == Never
+
+    t = A
+    assert eval_typing(GetArg[t, A, 0]) == Any
+    assert eval_typing(GetArg[t, A, -1]) == Any
+    assert eval_typing(GetArg[t, A, 1]) == Never
+
+
+def test_eval_getarg_custom_03():
+    class A[T = str]:
+        pass
+
+    t = A[int]
+    assert eval_typing(GetArg[t, A, 0]) is int
+    assert eval_typing(GetArg[t, A, -1]) is int
+    assert eval_typing(GetArg[t, A, 1]) == Never
+
+    t = A
+    assert eval_typing(GetArg[t, A, 0]) is str
+    assert eval_typing(GetArg[t, A, -1]) is str
+    assert eval_typing(GetArg[t, A, 1]) == Never
+
+
+def test_eval_getarg_custom_04():
+    T = TypeVar("T", default=str)
+
+    class A(Generic[T]):
+        pass
+
+    t = A[int]
+    assert eval_typing(GetArg[t, A, 0]) is int
+    assert eval_typing(GetArg[t, A, -1]) is int
+    assert eval_typing(GetArg[t, A, 1]) == Never
+
+    t = A
+    assert eval_typing(GetArg[t, A, 0]) is str
+    assert eval_typing(GetArg[t, A, -1]) is str
+    assert eval_typing(GetArg[t, A, 1]) == Never
+
+
 def test_uppercase_never():
     d = eval_typing(Uppercase[Never])
     assert d is Never

typemap/type_eval/_eval_operators.py

@@ -1,6 +1,10 @@
+import collections
+import collections.abc
+import contextlib
 import functools
 import inspect
 import itertools
+import re
 import types
 import typing
 
@@ -22,6 +26,7 @@
     Members,
     NewProtocol,
     Param,
+    SpecialFormEllipsis,
     StrConcat,
     StrSlice,
     Uncapitalize,
@@ -266,15 +271,125 @@ def _get_args(tp, base, ctx) -> typing.Any:
         return None
 
 
+def _fix_type(tp):
+    """Fix up a type getting returned from GetArg
+
+    In particular, this means turning a list into a tuple of the list
+    elements and turning ... into SpecialFormEllipsis.
+    """
+    if isinstance(tp, (tuple, list)):
+        return tuple[*tp]
+    elif tp is ...:
+        return SpecialFormEllipsis
+    else:
+        return tp
+
+
+# The number of generic parameters to all the builtin types that had
+# subscripting added in PEP 585.
+_BUILTIN_GENERIC_ARITIES = {
+    tuple: 2,  # variadic, like Callable...
+    list: 1,
+    dict: 2,
+    set: 1,
+    frozenset: 1,
+    type: 1,
+    collections.deque: 1,
+    collections.defaultdict: 2,
+    collections.OrderedDict: 2,
+    collections.Counter: 1,
+    collections.ChainMap: 2,
+    collections.abc.Awaitable: 1,
+    collections.abc.Coroutine: 3,
+    collections.abc.AsyncIterable: 1,
+    collections.abc.AsyncIterator: 1,
+    collections.abc.AsyncGenerator: 2,
+    collections.abc.Iterable: 1,
+    collections.abc.Iterator: 1,
+    collections.abc.Generator: 3,
+    collections.abc.Reversible: 1,
+    collections.abc.Container: 1,
+    collections.abc.Collection: 1,
+    collections.abc.Callable: 2,  # special syntax
+    collections.abc.Set: 1,
+    collections.abc.MutableSet: 1,
+    collections.abc.Mapping: 2,
+    collections.abc.MutableMapping: 2,
+    collections.abc.Sequence: 1,
+    collections.abc.MutableSequence: 1,
+    collections.abc.KeysView: 1,
+    collections.abc.ItemsView: 2,
+    collections.abc.ValuesView: 1,
+    contextlib.AbstractContextManager: 1,
+    contextlib.AbstractAsyncContextManager: 1,
+    re.Pattern: 1,
+    re.Match: 1,
+}
+
+
+def _get_params(base_head):
+    if (params := getattr(base_head, "__parameters__", None)) is not None:
+        return params
+    elif (params := getattr(base_head, "__type_params__", None)) is not None:
+        return params
+    else:
+        return None
+
+
+def _get_generic_arity(base_head):
+    if (n := _BUILTIN_GENERIC_ARITIES.get(base_head)) is not None:
+        return n
+    # XXX: check the type?
+    elif (n := getattr(base_head, "_nparams", None)) is not None:
+        return n
+    elif (params := _get_params(base_head)) is not None:
+        # TODO: also check for TypeVarTuple!
+        return len(params)
+    else:
+        return -1
+
+
+def _get_defaults(base_head):
+    """Get the *default* type params for a type
+
+    `list` is equivalent to `list[Any]`, so `GetArg[list, list, 0]
+    ought to return `Any`, while `GetArg[list, list, 1]` ought to
+    return `Never` because the index is invalid.
+
+    Annoyingly we need to consult a table for built-in arities for this.
+    """
+    arity = _get_generic_arity(base_head)
+    if arity < 0:
+        return None
+
+    # Callable and tuple need to produce a SpecialFormEllipsis for arg
+    # 0 and 1, respectively.
+    if base_head is collections.abc.Callable:
+        return (SpecialFormEllipsis, typing.Any)
+    elif base_head is tuple:
+        return (typing.Any, SpecialFormEllipsis)
+
+    if params := _get_params(base_head):
+        return tuple(
+            typing.Any if t.__default__ == typing.NoDefault else t.__default__
+            for t in params
+        )
+
+    return (typing.Any,) * arity
+
+
 @type_eval.register_evaluator(GetArg)
 @_lift_over_unions
 def _eval_GetArg(tp, base, idx, *, ctx) -> typing.Any:
-    args = _get_args(tp, base, ctx)
+    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
 
     try:
-        return args[_from_literal(idx, ctx)]
+        return _fix_type(args[_from_literal(idx, ctx)])
     except IndexError:
         return typing.Never