MAINT: simplify normalizer

Author: lezcano

torch_np/_normalizations.py

@@ -82,6 +82,36 @@ def normalize_ndarray(arg, name=None):
 import functools
 
 
+def maybe_normalize(arg, parm):
+    """Normalize arg if a normalizer is registred."""
+    normalizer = normalizers.get(parm.annotation, None)
+    return normalizer(arg) if normalizer else arg
+
+
+def normalizer(func):
+    @functools.wraps(func)
+    def wrapped(*args, **kwds):
+        params = inspect.signature(func).parameters
+        first_param = next(iter(params.values()))
+        # NumPy's API does not have positional args before variadic positional args
+        if first_param.kind == inspect.Parameter.VAR_POSITIONAL:
+            args = [maybe_normalize(arg, first_param) for arg in args]
+        else:
+            # NB: extra unknown arguments: pass through, will raise in func(*args) below
+            args = [
+                maybe_normalize(arg, parm) for arg, parm in zip(args, params.values())
+            ] + args[len(params.values()) :]
+
+        kwds = {
+            name: maybe_normalize(arg, params[name]) if name in params else arg
+            for name, arg in kwds.items()
+        }
+        return func(*args, **kwds)
+
+    return wrapped
+
+
+'''
 def normalize_this(arg, parm):
     """Normalize arg if a normalizer is registred."""
     normalizer = normalizers.get(parm.annotation, None)
@@ -123,3 +153,4 @@ def wrapped(*args, **kwds):
         return result
 
     return wrapped
+'''