ENH: annotate *args

This is a bit clumsy: func(*args : Annotation) gives a single annotation for a
runtime-determined number of arguments. There is no way to annotate individual
elements of *args AFAICS.

Thus register a special annotation to repack args into a tuple and a normalizer
to normalize this tuple.

Author: ev-br

torch_np/_funcs.py

@@ -13,6 +13,9 @@
 DTypeLike = typing.TypeVar("DTypeLike")
 SubokLike = typing.TypeVar("SubokLike")
 
+# annotate e.g. atleast_1d(*arys)
+UnpackedSeqArrayLike = typing.TypeVar("UnpackedSeqArrayLike")
+
 
 import inspect
 
@@ -53,30 +56,49 @@ def normalize_subok_like(arg, name):
     ArrayLike: normalize_array_like,
     Optional[ArrayLike]: normalize_optional_array_like,
     Sequence[ArrayLike]: normalize_seq_array_like,
+    UnpackedSeqArrayLike: normalize_seq_array_like,   # cf handling in normalize
     DTypeLike: normalize_dtype,
     SubokLike: normalize_subok_like,
 }
 
 import functools
 
 
+def normalize_this(arg, parm):
+    """Normalize arg if a normalizer is registred."""
+    normalizer = normalizers.get(parm.annotation, None)
+    if normalizer:
+        return normalizer(arg)
+    else:
+        # untyped arguments pass through
+        return arg
+
+
 def normalizer(func):
     @functools.wraps(func)
     def wrapped(*args, **kwds):
         sig = inspect.signature(func)
 
-        lst, dct = [], {}
+        # first, check for *args in positional parameters. Case in point:
+        # atleast_1d(*arys: UnpackedSequenceArrayLike)
+        # if found,  consume all args into a tuple to normalize as a whole
+        for j, param in enumerate(sig.parameters.values()):
+            if param.annotation == UnpackedSeqArrayLike:
+                if j == 0:
+                    args = (args,)
+                else:
+                    # args = args[:j] + (args[j:],) would likely work
+                    # not present in numpy codebase, so do not bother just yet.
+                    # NB: branching on j ==0 is to avoid the empty tuple, args[:j]
+                    raise NotImplementedError
+                break
+
         # loop over positional parameters and actual arguments
+        lst, dct = [], {}
         for arg, (name, parm) in zip(args, sig.parameters.items()):
             print(arg, name, parm.annotation)
-            normalizer = normalizers.get(parm.annotation, None)
-            if normalizer:
-                # dct[name] = normalizer(arg, name)
-                lst.append(normalizer(arg))
-            else:
-                # untyped arguments pass through
-                # dct[name] = arg
-                lst.append(arg)
+            lst.append(normalize_this(arg, parm))
+
 
         # normalize keyword arguments
         for name, arg in kwds.items():
@@ -88,11 +110,7 @@ def wrapped(*args, **kwds):
 
             print("kw: ", name, sig.parameters[name].annotation)
             parm = sig.parameters[name]
-            normalizer = normalizers.get(parm.annotation, None)
-            if normalizer:
-                dct[name] = normalizer(kwds[name], name)
-            else:
-                dct[name] = arg
+            dct[name] = normalize_this(arg, parm)
 
         ba = sig.bind(*lst, **dct)
         ba.apply_defaults()
@@ -113,6 +131,7 @@ def wrapped(*args, **kwds):
         # 5. axes : live in _impl or in types? several ways of handling them
         # 6. keepdims : peel off, postprocess
         # 7. OutLike : normal & keyword-only, peel off, postprocess
+        # 8. *args
 
         # finally, pass normalized arguments through
         result = func(*ba.args, **ba.kwargs)

torch_np/_wrapper.py

@@ -13,6 +13,11 @@
 
 from . import _dtypes, _helpers, _decorators  # isort: skip  # XXX
 
+### XXX: order the imports DAG
+from . _funcs import normalizer, DTypeLike, ArrayLike, UnpackedSeqArrayLike
+from typing import Optional, Sequence
+
+
 # Things to decide on (punt for now)
 #
 # 1. Q: What are the return types of wrapper functions: plain torch.Tensors or
@@ -63,27 +68,27 @@ def copy(a, order="K", subok=False):
     return a.copy(order="C")
 
 
-def atleast_1d(*arys):
-    tensors = _helpers.to_tensors(*arys)
-    res = torch.atleast_1d(tensors)
+@normalizer
+def atleast_1d(*arys : UnpackedSeqArrayLike):
+    res = torch.atleast_1d(*arys)
     if len(res) == 1:
         return asarray(res[0])
     else:
         return list(asarray(_) for _ in res)
 
 
-def atleast_2d(*arys):
-    tensors = _helpers.to_tensors(*arys)
-    res = torch.atleast_2d(tensors)
+@normalizer
+def atleast_2d(*arys : UnpackedSeqArrayLike):
+    res = torch.atleast_2d(*arys)
     if len(res) == 1:
         return asarray(res[0])
     else:
         return list(asarray(_) for _ in res)
 
 
-def atleast_3d(*arys):
-    tensors = _helpers.to_tensors(*arys)
-    res = torch.atleast_3d(tensors)
+@normalizer
+def atleast_3d(*arys : UnpackedSeqArrayLike):
+    res = torch.atleast_3d(*arys)
     if len(res) == 1:
         return asarray(res[0])
     else:
@@ -108,10 +113,6 @@ def _concat_check(tup, dtype, out):
             )
 
 
-### XXX: order the imports DAG
-from . _funcs import normalizer, DTypeLike, ArrayLike
-from typing import Optional, Sequence
-
 @normalizer
 def concatenate(ar_tuple : Sequence[ArrayLike], axis=0, out=None, dtype: DTypeLike=None, casting="same_kind"):
     _concat_check(ar_tuple, dtype, out=out)