@@ -164,7 +164,8 @@ PyCode_NewWithPosOnlyArgs(int argcount, int posonlyargcount, int kwonlyargcount,
164164{
165165 PyCodeObject * co ;
166166 Py_ssize_t * cell2arg = NULL ;
167- Py_ssize_t i , n_cellvars , n_varnames , total_args ;
167+ Py_ssize_t i , total_args ;
168+ int ncellvars , nfreevars ;
168169
169170 /* Check argument types */
170171 if (argcount < posonlyargcount || posonlyargcount < 0 ||
@@ -184,6 +185,19 @@ PyCode_NewWithPosOnlyArgs(int argcount, int posonlyargcount, int kwonlyargcount,
184185 return NULL ;
185186 }
186187
188+ /* Make sure that code is indexable with an int, this is
189+ a long running assumption in ceval.c and many parts of
190+ the interpreter. */
191+ if (PyBytes_GET_SIZE (code ) > INT_MAX ) {
192+ PyErr_SetString (PyExc_OverflowError , "co_code larger than INT_MAX" );
193+ return NULL ;
194+ }
195+
196+ if (nlocals != PyTuple_GET_SIZE (varnames )) {
197+ PyErr_SetString (PyExc_ValueError , "co_nlocals != len(co_varnames)" );
198+ return NULL ;
199+ }
200+
187201 /* Ensure that strings are ready Unicode string */
188202 if (PyUnicode_READY (name ) < 0 ) {
189203 return NULL ;
@@ -208,46 +222,40 @@ PyCode_NewWithPosOnlyArgs(int argcount, int posonlyargcount, int kwonlyargcount,
208222 return NULL ;
209223 }
210224
211- /* Make sure that code is indexable with an int, this is
212- a long running assumption in ceval.c and many parts of
213- the interpreter. */
214- if (PyBytes_GET_SIZE (code ) > INT_MAX ) {
215- PyErr_SetString (PyExc_OverflowError , "co_code larger than INT_MAX" );
216- return NULL ;
217- }
218-
219225 /* Check for any inner or outer closure references */
220- n_cellvars = PyTuple_GET_SIZE (cellvars );
221- if (!n_cellvars && !PyTuple_GET_SIZE (freevars )) {
226+ assert (PyTuple_GET_SIZE (cellvars ) < INT_MAX );
227+ ncellvars = (int )PyTuple_GET_SIZE (cellvars );
228+ assert (PyTuple_GET_SIZE (freevars ) < INT_MAX );
229+ nfreevars = (int )PyTuple_GET_SIZE (freevars );
230+ if (!ncellvars && !nfreevars ) {
222231 flags |= CO_NOFREE ;
223232 } else {
224233 flags &= ~CO_NOFREE ;
225234 }
226235
227- n_varnames = PyTuple_GET_SIZE (varnames );
228- if (argcount <= n_varnames && kwonlyargcount <= n_varnames ) {
236+ if (argcount <= nlocals && kwonlyargcount <= nlocals ) {
229237 /* Never overflows. */
230238 total_args = (Py_ssize_t )argcount + (Py_ssize_t )kwonlyargcount +
231239 ((flags & CO_VARARGS ) != 0 ) + ((flags & CO_VARKEYWORDS ) != 0 );
232240 }
233241 else {
234- total_args = n_varnames + 1 ;
242+ total_args = nlocals + 1 ;
235243 }
236- if (total_args > n_varnames ) {
244+ if (total_args > nlocals ) {
237245 PyErr_SetString (PyExc_ValueError , "code: varnames is too small" );
238246 return NULL ;
239247 }
240248
241249 /* Create mapping between cells and arguments if needed. */
242- if (n_cellvars ) {
250+ if (ncellvars ) {
243251 bool used_cell2arg = false;
244- cell2arg = PyMem_NEW (Py_ssize_t , n_cellvars );
252+ cell2arg = PyMem_NEW (Py_ssize_t , ncellvars );
245253 if (cell2arg == NULL ) {
246254 PyErr_NoMemory ();
247255 return NULL ;
248256 }
249257 /* Find cells which are also arguments. */
250- for (i = 0 ; i < n_cellvars ; i ++ ) {
258+ for (i = 0 ; i < ncellvars ; i ++ ) {
251259 Py_ssize_t j ;
252260 PyObject * cell = PyTuple_GET_ITEM (cellvars , i );
253261 cell2arg [i ] = CO_CELL_NOT_AN_ARG ;
@@ -279,9 +287,10 @@ PyCode_NewWithPosOnlyArgs(int argcount, int posonlyargcount, int kwonlyargcount,
279287 co -> co_argcount = argcount ;
280288 co -> co_posonlyargcount = posonlyargcount ;
281289 co -> co_kwonlyargcount = kwonlyargcount ;
290+ co -> co_nlocalsplus = nlocals + ncellvars + nfreevars ;
282291 co -> co_nlocals = nlocals ;
283- co -> co_nlocalsplus = nlocals +
284- ( int ) PyTuple_GET_SIZE ( freevars ) + ( int ) PyTuple_GET_SIZE ( cellvars ) ;
292+ co -> co_ncellvars = ncellvars ;
293+ co -> co_nfreevars = nfreevars ;
285294 co -> co_stacksize = stacksize ;
286295 co -> co_flags = flags ;
287296 Py_INCREF (code );
@@ -1139,7 +1148,7 @@ code_sizeof(PyCodeObject *co, PyObject *Py_UNUSED(args))
11391148 _PyCodeObjectExtra * co_extra = (_PyCodeObjectExtra * ) co -> co_extra ;
11401149
11411150 if (co -> co_cell2arg != NULL && co -> co_cellvars != NULL ) {
1142- res += PyTuple_GET_SIZE ( co -> co_cellvars ) * sizeof (Py_ssize_t );
1151+ res += co -> co_ncellvars * sizeof (Py_ssize_t );
11431152 }
11441153 if (co_extra != NULL ) {
11451154 res += sizeof (_PyCodeObjectExtra ) +
0 commit comments