Improve compatibility with "nogil" Python and 3.11 (#470)

* Improve compatibility with "nogil" Python and 3.11

This makes a number of changes to improve compatibility with the
"nogil" Python fork as well as the upcoming 3.11 release.

 - Fix _code_reduce for 3.11b0 and nogil Python

 - Use instr.argval in _walk_global_ops. This avoids adding a special
   case for 3.11+ (and is useful for nogil Python). In 3.11+, the argval
   for LOAD_GLOBAL would need to be divided by two to access the correct
   name. The 'argval' field already stores the correct name.

 - Set '__builtins__' before constructing de-pickled functions. (Useful
   for nogil Python)

 - Fix test_recursion_during_pickling in Python 3.11+. Objects now have
   a default `__getstate__` method so `__getattr__` was never called,
   but `__getattribute__` would still be called.


Co-authored-by: Olivier Grisel <[email protected]>

CHANGES.md

@@ -11,6 +11,9 @@
 - Support and CI configuration for Python 3.11.
   ([PR #467](https://github.com/cloudpipe/cloudpickle/pull/467))
 
+- Support for the experimental `nogil` variant of CPython
+  ([PR #470](https://github.com/cloudpipe/cloudpickle/pull/470))
+
 2.0.0
 =====
 

cloudpickle/cloudpickle.py

@@ -321,11 +321,10 @@ def _extract_code_globals(co):
     """
     out_names = _extract_code_globals_cache.get(co)
     if out_names is None:
-        names = co.co_names
         # We use a dict with None values instead of a set to get a
         # deterministic order (assuming Python 3.6+) and avoid introducing
         # non-deterministic pickle bytes as a results.
-        out_names = {names[oparg]: None for _, oparg in _walk_global_ops(co)}
+        out_names = {name: None for name in _walk_global_ops(co)}
 
         # Declaring a function inside another one using the "def ..."
         # syntax generates a constant code object corresponding to the one
@@ -511,13 +510,12 @@ def _builtin_type(name):
 
 def _walk_global_ops(code):
     """
-    Yield (opcode, argument number) tuples for all
-    global-referencing instructions in *code*.
+    Yield referenced name for all global-referencing instructions in *code*.
     """
     for instr in dis.get_instructions(code):
         op = instr.opcode
         if op in GLOBAL_OPS:
-            yield op, instr.arg
+            yield instr.argval
 
 
 def _extract_class_dict(cls):
@@ -765,6 +763,12 @@ def _fill_function(*args):
     return func
 
 
+def _make_function(code, globals, name, argdefs, closure):
+    # Setting __builtins__ in globals is needed for nogil CPython.
+    globals["__builtins__"] = __builtins__
+    return types.FunctionType(code, globals, name, argdefs, closure)
+
+
 def _make_empty_cell():
     if False:
         # trick the compiler into creating an empty cell in our lambda

cloudpickle/cloudpickle_fast.py

@@ -35,7 +35,7 @@
     _is_parametrized_type_hint, PYPY, cell_set,
     parametrized_type_hint_getinitargs, _create_parametrized_type_hint,
     builtin_code_type,
-    _make_dict_keys, _make_dict_values, _make_dict_items,
+    _make_dict_keys, _make_dict_values, _make_dict_items, _make_function,
 )
 
 
@@ -248,17 +248,16 @@ def _code_reduce(obj):
     # of the specific type from types, for example:
     # >>> from types import CodeType
     # >>> help(CodeType)
-    if hasattr(obj, "co_columntable"):  # pragma: no branch
+    if hasattr(obj, "co_exceptiontable"):  # pragma: no branch
         # Python 3.11 and later: there are some new attributes
         # related to the enhanced exceptions.
         args = (
             obj.co_argcount, obj.co_posonlyargcount,
             obj.co_kwonlyargcount, obj.co_nlocals, obj.co_stacksize,
             obj.co_flags, obj.co_code, obj.co_consts, obj.co_names,
             obj.co_varnames, obj.co_filename, obj.co_name, obj.co_qualname,
-            obj.co_firstlineno, obj.co_linetable, obj.co_endlinetable,
-            obj.co_columntable, obj.co_exceptiontable, obj.co_freevars,
-            obj.co_cellvars,
+            obj.co_firstlineno, obj.co_linetable, obj.co_exceptiontable,
+            obj.co_freevars, obj.co_cellvars,
         )
     elif hasattr(obj, "co_linetable"):  # pragma: no branch
         # Python 3.10 and later: obj.co_lnotab is deprecated and constructor
@@ -271,6 +270,18 @@ def _code_reduce(obj):
             obj.co_firstlineno, obj.co_linetable, obj.co_freevars,
             obj.co_cellvars
         )
+    elif hasattr(obj, "co_nmeta"):  # pragma: no cover
+        # "nogil" Python: modified attributes from 3.9
+        args = (
+            obj.co_argcount, obj.co_posonlyargcount,
+            obj.co_kwonlyargcount, obj.co_nlocals, obj.co_framesize,
+            obj.co_ndefaultargs, obj.co_nmeta,
+            obj.co_flags, obj.co_code, obj.co_consts,
+            obj.co_varnames, obj.co_filename, obj.co_name,
+            obj.co_firstlineno, obj.co_lnotab, obj.co_exc_handlers,
+            obj.co_jump_table, obj.co_freevars, obj.co_cellvars,
+            obj.co_free2reg, obj.co_cell2reg
+        )
     elif hasattr(obj, "co_posonlyargcount"):
         # Backward compat for 3.9 and older
         args = (
@@ -564,7 +575,7 @@ def _dynamic_function_reduce(self, func):
         """Reduce a function that is not pickleable via attribute lookup."""
         newargs = self._function_getnewargs(func)
         state = _function_getstate(func)
-        return (types.FunctionType, newargs, state, None, None,
+        return (_make_function, newargs, state, None, None,
                 _function_setstate)
 
     def _function_reduce(self, obj):

tests/cloudpickle_test.py

@@ -978,6 +978,10 @@ def g(y):
         res = loop.run_sync(functools.partial(g2, 5))
         self.assertEqual(res, 7)
 
+    @pytest.mark.skipif(
+        (3, 11, 0, 'beta') <= sys.version_info < (3, 11, 0, 'beta', 2),
+        reason="https://github.com/python/cpython/issues/92932"
+    )
     def test_extended_arg(self):
         # Functions with more than 65535 global vars prefix some global
         # variable references with the EXTENDED_ARG opcode.
@@ -2245,7 +2249,7 @@ def inner_function():
 
     def test_recursion_during_pickling(self):
         class A:
-            def __getattr__(self, name):
+            def __getattribute__(self, name):
                 return getattr(self, name)
 
         a = A()