Move the C code to Python.

python · serhiy-storchaka · Jun 12, 2022 · May 29, 2022 · May 29, 2022 · May 29, 2022
commit a7efb9135ad96180e56cc1bb51c896b4edd40d30
diff --git a/Include/internal/pycore_global_strings.h b/Include/internal/pycore_global_strings.h
@@ -202,6 +202,7 @@ struct _Py_global_strings {
         STRUCT_FOR_ID(__truediv__)
         STRUCT_FOR_ID(__trunc__)
         STRUCT_FOR_ID(__typing_is_unpacked_typevartuple__)
+        STRUCT_FOR_ID(__typing_prepare_subst__)
         STRUCT_FOR_ID(__typing_subst__)
         STRUCT_FOR_ID(__typing_unpacked_tuple_args__)
         STRUCT_FOR_ID(__warningregistry__)

@@ -824,6 +824,7 @@ extern "C" {
                 INIT_ID(__truediv__), \
                 INIT_ID(__trunc__), \
                 INIT_ID(__typing_is_unpacked_typevartuple__), \
+                INIT_ID(__typing_prepare_subst__), \
                 INIT_ID(__typing_subst__), \
                 INIT_ID(__typing_unpacked_tuple_args__), \
                 INIT_ID(__warningregistry__), \

@@ -1065,6 +1065,42 @@ def __repr__(self):
     def __typing_subst__(self, arg):
         raise TypeError("Substitution of bare TypeVarTuple is not supported")
 
+    def __typing_prepare_subst__(self, alias, args):
+        params = alias.__parameters__
+        typevartuple_index = params.index(self)
+        for param in enumerate(params[typevartuple_index + 1:]):
+            if isinstance(param, TypeVarTuple):
+                raise TypeError(f"More than one TypeVarTuple parameter in {alias}")
+
+        alen = len(args)
+        plen = len(params)
+        left = typevartuple_index
+        right = plen - typevartuple_index - 1
+        var_tuple_index = None
+        fillarg = None
+        for k, arg in enumerate(args):
+            if not (isinstance(arg, type) and not isinstance(arg, GenericAlias)):
+                subargs = getattr(arg, '__typing_unpacked_tuple_args__', None)
+                if subargs and len(subargs) == 2 and subargs[-1] is ...:
+                    if var_tuple_index is not None:
+                        raise TypeError("More than one unpacked arbitrary-length tuple argument")
+                    var_tuple_index = k
+                    fillarg = subargs[0]
+        if var_tuple_index is not None:
+            left = min(left, var_tuple_index)
+            right = min(right, alen - var_tuple_index - 1)
+        elif left + right > alen:
+            raise TypeError(f"Too few arguments for {alias};"
+                            f" actual {alen}, expected at least {plen-1}")
+
+        return (
+            *args[:left],
+            *([fillarg]*(typevartuple_index - left)),
+            tuple(args[left: alen - right]),
+            *([fillarg]*(plen - right - left - typevartuple_index - 1)),
+            *args[alen - right:],
+        )
+
 
 class ParamSpecArgs(_Final, _Immutable, _root=True):
     """The args for a ParamSpec object.
@@ -1184,6 +1220,8 @@ def __typing_subst__(self, arg):
                             f"ParamSpec, or Concatenate. Got {arg}")
         return arg
 
+    def __typing_prepare_subst__(self, alias, args):
+        return _prepare_paramspec_params(alias, args)
 
 def _is_dunder(attr):
     return attr.startswith('__') and attr.endswith('__')
@@ -1347,10 +1385,6 @@ def __getitem__(self, args):
             args = (args,)
         args = tuple(_type_convert(p) for p in args)
         args = _unpack_args(args)
-        if (self._paramspec_tvars
-                and any(isinstance(t, ParamSpec) for t in self.__parameters__)):
-            args = _prepare_paramspec_params(self, args)
-
         new_args = self._determine_new_args(args)
         r = self.copy_with(new_args)
         return r
@@ -1372,47 +1406,16 @@ def _determine_new_args(self, args):
 
         params = self.__parameters__
         # In the example above, this would be {T3: str}
-        new_arg_by_param = {}
-        typevartuple_index = None
-        for i, param in enumerate(params):
-            if isinstance(param, TypeVarTuple):
-                if typevartuple_index is not None:
-                    raise TypeError(f"More than one TypeVarTuple parameter in {self}")
-                typevartuple_index = i
-
+        for param in params:
+            prepare = getattr(param, '__typing_prepare_subst__', None)
+            if prepare is not None:
+                args = prepare(self, args)
         alen = len(args)
         plen = len(params)
-        if typevartuple_index is not None:
-            left = typevartuple_index
-            right = plen - typevartuple_index - 1
-            var_tuple_index = None
-            for k, arg in enumerate(args):
-                if not (isinstance(arg, type) and not isinstance(arg, GenericAlias)):
-                    subargs = getattr(arg, '__typing_unpacked_tuple_args__', None)
-                    if subargs and len(subargs) == 2 and subargs[-1] is ...:
-                        if var_tuple_index is not None:
-                            raise TypeError("More than one unpacked arbitrary-length tuple argument")
-                        var_tuple_index = k
-                        fillarg = subargs[0]
-            if var_tuple_index is not None:
-                left = min(left, var_tuple_index)
-                right = min(right, alen - var_tuple_index - 1)
-            elif left + right > alen:
-                raise TypeError(f"Too few arguments for {self};"
-                                f" actual {alen}, expected at least {plen-1}")
-
-            new_arg_by_param.update(zip(params[:left], args[:left]))
-            for k in range(left, typevartuple_index):
-                new_arg_by_param[params[k]] = fillarg
-            new_arg_by_param[params[typevartuple_index]] = tuple(args[left: alen - right])
-            for k in range(typevartuple_index + 1, plen - right):
-                new_arg_by_param[params[k]] = fillarg
-            new_arg_by_param.update(zip(params[plen - right:], args[alen - right:]))
-        else:
-            if alen != plen:
-                raise TypeError(f"Too {'many' if alen > plen else 'few'} arguments for {self};"
-                                f" actual {alen}, expected {plen}")
-            new_arg_by_param.update(zip(params, args))
+        if alen != plen:
+            raise TypeError(f"Too {'many' if alen > plen else 'few'} arguments for {self};"
+                            f" actual {alen}, expected {plen}")
+        new_arg_by_param = dict(zip(params, args))
 
         new_args = []
         for old_arg in self.__args__:

diff --git a/Objects/genericaliasobject.c b/Objects/genericaliasobject.c
@@ -269,9 +269,7 @@ _Py_make_parameters(PyObject *args)
    a non-empty tuple, return a new reference to obj. */
 static PyObject *
 subs_tvars(PyObject *obj, PyObject *params,
-           PyObject **argitems, Py_ssize_t nargs,
-           Py_ssize_t varparam, Py_ssize_t left, Py_ssize_t right,
-           PyObject *fillarg)
+           PyObject **argitems, Py_ssize_t nargs)
 {
     PyObject *subparams;
     if (_PyObject_LookupAttr(obj, &_Py_ID(__parameters__), &subparams) < 0) {
@@ -289,30 +287,22 @@ subs_tvars(PyObject *obj, PyObject *params,
         for (Py_ssize_t i = 0; i < nsubargs; ++i) {
             PyObject *arg = PyTuple_GET_ITEM(subparams, i);
             Py_ssize_t iparam = tuple_index(params, nparams, arg);
-            if (iparam == varparam) {
-                j = tuple_extend(&subargs, j,
-                                 argitems + left, nargs - left - right);
-                if (j < 0) {
-                    return NULL;
-                }
-            }
-            else {
-                if (iparam >= 0) {
-                    if (iparam < left) {
-                        arg = argitems[iparam];
-                    }
-                    else if (iparam >= nparams - right) {
-                        iparam += nargs - nparams;
-                        arg = argitems[iparam];
-                    }
-                    else {
-                        arg = fillarg;
+            if (iparam >= 0) {
+                PyObject *param = PyTuple_GET_ITEM(params, iparam);
+                arg = argitems[iparam];
+                if (Py_TYPE(param)->tp_iter && PyTuple_Check(arg)) {  // TypeVarTuple
+                    j = tuple_extend(&subargs, j,
+                                    &PyTuple_GET_ITEM(arg, 0),
+                                    PyTuple_GET_SIZE(arg));
+                    if (j < 0) {
+                        return NULL;
                     }
+                    continue;
                 }
-                Py_INCREF(arg);
-                PyTuple_SET_ITEM(subargs, j, arg);
-                j++;
             }
+            Py_INCREF(arg);
+            PyTuple_SET_ITEM(subargs, j, arg);
+            j++;
         }
         assert(j == PyTuple_GET_SIZE(subargs));
 
@@ -409,27 +399,6 @@ _unpack_args(PyObject *item)
     return newargs;
 }
 
-static PyObject *
-_get_unpacked_var_tuple_arg(PyObject *arg)
-{
-    if (PyType_Check(arg)) {
-        return NULL;
-    }
-    PyObject *subargs = _unpacked_tuple_args(arg);
-    if (subargs != NULL &&
-        PyTuple_Check(subargs) &&
-        PyTuple_GET_SIZE(subargs) == 2 &&
-        PyTuple_GET_ITEM(subargs, 1) == Py_Ellipsis)
-    {
-        PyObject *subarg = PyTuple_GET_ITEM(subargs, 0);
-        Py_INCREF(subarg);
-        Py_DECREF(subargs);
-        return subarg;
-    }
-    Py_XDECREF(subargs);
-    return NULL;
-}
-
 PyObject *
 _Py_subs_parameters(PyObject *self, PyObject *args, PyObject *parameters, PyObject *item)
 {
@@ -440,67 +409,36 @@ _Py_subs_parameters(PyObject *self, PyObject *args, PyObject *parameters, PyObje
                             self);
     }
     item = _unpack_args(item);
-    int is_tuple = PyTuple_Check(item);
-    Py_ssize_t nitems = is_tuple ? PyTuple_GET_SIZE(item) : 1;
-    PyObject **argitems = is_tuple ? &PyTuple_GET_ITEM(item, 0) : &item;
-    Py_ssize_t varparam = nparams;
     for (Py_ssize_t i = 0; i < nparams; i++) {
         PyObject *param = PyTuple_GET_ITEM(parameters, i);
-        if (Py_TYPE(param)->tp_iter) {  // TypeVarTuple
-            if (varparam < nparams) {
-                Py_DECREF(item);
-                return PyErr_Format(PyExc_TypeError,
-                                    "More than one TypeVarTuple parameter in %S",
-                                    self);
-            }
-            varparam = i;
+        PyObject *prepare, *tmp;
+        if (_PyObject_LookupAttr(param, &_Py_ID(__typing_prepare_subst__), &prepare) < 0) {
+            Py_DECREF(item);
+            return NULL;
         }
-    }
-    PyObject *fillarg = NULL;
-    Py_ssize_t vartuplearg = nitems;
-    Py_ssize_t left = varparam;
-    Py_ssize_t right = nparams - varparam - 1;
-    if (varparam < nparams) {
-        for (Py_ssize_t i = 0; i < nitems; i++) {
-            PyObject *arg = _get_unpacked_var_tuple_arg(argitems[i]);
-            if (arg) {
-                if (vartuplearg < nitems) {
-                    Py_DECREF(arg);
-                    Py_DECREF(fillarg);
-                    Py_DECREF(item);
-                    return PyErr_Format(PyExc_TypeError,
-                                        "More than one unpacked arbitrary-length tuple argument",
-                                        self);
-                }
-                vartuplearg = i;
-                fillarg = arg;
+        if (prepare && prepare != Py_None) {
+            if (PyTuple_Check(item)) {
+                tmp = PyObject_CallFunction(prepare, "OO", self, item);
             }
-            else if (PyErr_Occurred()) {
-                Py_XDECREF(fillarg);
-                Py_DECREF(item);
+            else {
+                tmp = PyObject_CallFunction(prepare, "O(O)", self, item);
+            }
+            Py_DECREF(prepare);
+            Py_SETREF(item, tmp);
+            if (item == NULL) {
                 return NULL;
             }
         }
-        if (vartuplearg < nitems) {
-            assert(fillarg);
-            left = Py_MIN(left, vartuplearg);
-            right = Py_MIN(right, nitems - vartuplearg - 1);
-        }
-        else if (left + right > nitems) {
-            Py_DECREF(item);
-            return PyErr_Format(PyExc_TypeError,
-                                "Too few arguments for %R",
-                                self);
-        }
     }
-    else {
-        if (nitems != nparams) {
-            Py_DECREF(item);
-            return PyErr_Format(PyExc_TypeError,
-                                "Too %s arguments for %R; actual %zd, expected %zd",
-                                nitems > nparams ? "many" : "few",
-                                self, nitems, nparams);
-        }
+    int is_tuple = PyTuple_Check(item);
+    Py_ssize_t nitems = is_tuple ? PyTuple_GET_SIZE(item) : 1;
+    PyObject **argitems = is_tuple ? &PyTuple_GET_ITEM(item, 0) : &item;
+    if (nitems != nparams) {
+        Py_DECREF(item);
+        return PyErr_Format(PyExc_TypeError,
+                            "Too %s arguments for %R; actual %zd, expected %zd",
+                            nitems > nparams ? "many" : "few",
+                            self, nitems, nparams);
     }
     /* Replace all type variables (specified by parameters)
        with corresponding values specified by argitems.
@@ -511,7 +449,6 @@ _Py_subs_parameters(PyObject *self, PyObject *args, PyObject *parameters, PyObje
     Py_ssize_t nargs = PyTuple_GET_SIZE(args);
     PyObject *newargs = PyTuple_New(nargs);
     if (newargs == NULL) {
-        Py_XDECREF(fillarg);
         Py_DECREF(item);
         return NULL;
     }
@@ -520,50 +457,26 @@ _Py_subs_parameters(PyObject *self, PyObject *args, PyObject *parameters, PyObje
         int unpack = _is_unpacked_typevartuple(arg);
         if (unpack < 0) {
             Py_DECREF(newargs);
-            Py_XDECREF(fillarg);
             Py_DECREF(item);
             return NULL;
         }
         PyObject *subst;
         if (_PyObject_LookupAttr(arg, &_Py_ID(__typing_subst__), &subst) < 0) {
             Py_DECREF(newargs);
-            Py_XDECREF(fillarg);
             Py_DECREF(item);
             return NULL;
         }
         if (subst) {
             Py_ssize_t iparam = tuple_index(parameters, nparams, arg);
             assert(iparam >= 0);
-            if (iparam == varparam) {
-                Py_DECREF(subst);
-                Py_DECREF(newargs);
-                Py_XDECREF(fillarg);
-                Py_DECREF(item);
-                PyErr_SetString(PyExc_TypeError,
-                        "Substitution of bare TypeVarTuple is not supported");
-                return NULL;
-            }
-            if (iparam < left) {
-                arg = argitems[iparam];
-            }
-            else if (iparam >= nparams - right) {
-                iparam += nitems - nparams;
-                arg = argitems[iparam];
-            }
-            else {
-                assert(fillarg);
-                arg = fillarg;
-            }
-            arg = PyObject_CallOneArg(subst, arg);
+            arg = PyObject_CallOneArg(subst, argitems[iparam]);
             Py_DECREF(subst);
         }
         else {
-            arg = subs_tvars(arg, parameters, argitems, nitems,
-                             varparam, left, right, fillarg);
+            arg = subs_tvars(arg, parameters, argitems, nitems);
         }
         if (arg == NULL) {
             Py_DECREF(newargs);
-            Py_XDECREF(fillarg);
             Py_DECREF(item);
             return NULL;
         }
@@ -572,7 +485,6 @@ _Py_subs_parameters(PyObject *self, PyObject *args, PyObject *parameters, PyObje
                     &PyTuple_GET_ITEM(arg, 0), PyTuple_GET_SIZE(arg));
             Py_DECREF(arg);
             if (jarg < 0) {
-                Py_XDECREF(fillarg);
                 Py_DECREF(item);
                 return NULL;
             }
@@ -583,7 +495,6 @@ _Py_subs_parameters(PyObject *self, PyObject *args, PyObject *parameters, PyObje
         }
     }
 
-    Py_XDECREF(fillarg);
     Py_DECREF(item);
     return newargs;
 }