Make cpdef enums pickleable (GH-5124)

(on Python >= 3.4). Closes #5120
cython · Jan 10, 2023 · ac56a66 · ac56a66
1 parent f58bfdb
commit ac56a66
Show file tree

Hide file tree

Showing 4 changed files with 36 additions and 3 deletions.
diff --git a/Cython/Compiler/PyrexTypes.py b/Cython/Compiler/PyrexTypes.py
@@ -4277,6 +4277,7 @@ def create_type_wrapper(self, env):
                 "items": tuple(self.values),
                 "underlying_type": self.underlying_type.empty_declaration_code(),
                 "enum_doc": self.doc,
+                "static_modname": env.qualified_name,
             },
             outer_module_scope=env.global_scope())
 
@@ -4387,6 +4388,7 @@ def create_type_wrapper(self, env):
             context={"name": self.name,
                      "items": tuple(self.values),
                      "enum_doc": self.doc,
+                     "static_modname": env.qualified_name,
                      },
             outer_module_scope=env.global_scope()))
 

diff --git a/Cython/Utility/CpdefEnums.pyx b/Cython/Utility/CpdefEnums.pyx
@@ -54,7 +54,8 @@ if PY_VERSION_HEX >= 0x03040000:
         {{for item in items}}
         ('{{item}}', {{item}}),
         {{endfor}}
-    ])
+        # Try to look up the module name dynamically if possible
+    ], module=__Pyx_globals.get("__module__", '{{static_modname}}'))
     {{if enum_doc is not None}}
     {{name}}.__doc__ = {{ repr(enum_doc) }}
     {{endif}}
@@ -79,8 +80,7 @@ if PY_VERSION_HEX >= 0x03040000:
         {{for item in items}}
         ('{{item}}', <{{underlying_type}}>({{name}}.{{item}})),
         {{endfor}}
-    ])
-
+    ], module=__Pyx_globals.get("__module__", '{{static_modname}}'))
 else:
     __Pyx_globals["{{name}}"] = type('{{name}}', (__Pyx_EnumBase,), {})
     {{for item in items}}

diff --git a/tests/run/cpdef_enums.pyx b/tests/run/cpdef_enums.pyx
@@ -131,3 +131,18 @@ def check_docs():
     'Home is where...'
     """
     pass
+
+def test_pickle():
+    """
+    >>> from pickle import loads, dumps
+    >>> import sys
+
+    Pickling enums won't work without the enum module, so disable the test
+    >>> if sys.version_info < (3, 4):
+    ...     loads = dumps = lambda x: x
+    >>> loads(dumps(PyxEnum.TWO)) == PyxEnum.TWO
+    True
+    >>> loads(dumps(PxdEnum.RANK_2)) == PxdEnum.RANK_2
+    True
+    """
+    pass
diff --git a/tests/run/cpdef_scoped_enums.pyx b/tests/run/cpdef_scoped_enums.pyx
@@ -40,3 +40,19 @@ def test_enum_doc():
     True
     """
     pass
+
+
+def test_pickle():
+    """
+    >>> from pickle import loads, dumps
+    >>> import sys
+
+    Pickling enums won't work without the enum module, so disable the test
+    >>> if sys.version_info < (3, 4):
+    ...     loads = dumps = lambda x: x
+    >>> loads(dumps(Enum1.Item2)) == Enum1.Item2
+    True
+    >>> loads(dumps(Enum2.Item4)) == Enum2.Item4
+    True
+    """
+    pass