fixes to support pickling of enums (by Shashank Bapat)

Ralf W. Grosse-Kunstleve · Ralf W. Grosse-Kunstleve · commit d67b0406834d · 2006-12-03T20:43:48.000Z
[SVN r36256]
diff --git a/src/object/enum.cpp b/src/object/enum.cpp
@@ -32,18 +32,19 @@ extern "C"
 {
     static PyObject* enum_repr(PyObject* self_)
     {
+        const char *mod = PyString_AsString(PyObject_GetAttrString( self_, "__module__"));
         enum_object* self = downcast<enum_object>(self_);
         if (!self->name)
         {
-            return PyString_FromFormat("%s(%ld)", self_->ob_type->tp_name, PyInt_AS_LONG(self_));
+            return PyString_FromFormat("%s.%s(%ld)", mod, self_->ob_type->tp_name, PyInt_AS_LONG(self_));
         }
         else
         {
             char* name = PyString_AsString(self->name);
             if (name == 0)
                 return 0;
             
-            return PyString_FromFormat("%s.%s", self_->ob_type->tp_name, name);
+            return PyString_FromFormat("%s.%s.%s", mod, self_->ob_type->tp_name, name);
         }
     }
 
@@ -141,10 +142,9 @@ namespace
 
       object module_name = module_prefix();
       if (module_name)
-          module_name += '.';
+         d["__module__"] = module_name;
       
-      object result = (object(metatype))(
-          module_name + name, make_tuple(base), d);
+      object result = (object(metatype))(name, make_tuple(base), d);
       
       scope().attr(name) = result;
 
diff --git a/test/enum.py b/test/enum.py
@@ -48,14 +48,27 @@
 enum_ext.color.green
 '''
 
+# pickling of enums only works with Python 2.3 or higher
+exercise_pickling = '''
+>>> import pickle
+>>> p = pickle.dumps(color.green, pickle.HIGHEST_PROTOCOL)
+>>> l = pickle.loads(p)
+>>> identity(l)
+enum_ext.color.green
+'''
+
 def run(args = None):
     import sys
     import doctest
+    import pickle
 
     if args is not None:
         sys.argv = args
-    return doctest.testmod(sys.modules.get(__name__))
-    
+    self = sys.modules.get(__name__)
+    if (hasattr(pickle, "HIGHEST_PROTOCOL")):
+      self.__doc__ += exercise_pickling
+    return doctest.testmod(self)
+
 if __name__ == '__main__':
     print "running..."
     import sys

Original file line number	Diff line number	Diff line change
`@@ -32,18 +32,19 @@ extern "C"`
`32`	`32`	`{`
`33`	`33`	`static PyObject* enum_repr(PyObject* self_)`
`34`	`34`	`{`
	`35`	`+ const char *mod = PyString_AsString(PyObject_GetAttrString( self_, "__module__"));`
`35`	`36`	`enum_object* self = downcast<enum_object>(self_);`
`36`	`37`	`if (!self->name)`
`37`	`38`	`{`
`38`		`- return PyString_FromFormat("%s(%ld)", self_->ob_type->tp_name, PyInt_AS_LONG(self_));`
	`39`	`+ return PyString_FromFormat("%s.%s(%ld)", mod, self_->ob_type->tp_name, PyInt_AS_LONG(self_));`
`39`	`40`	`}`
`40`	`41`	`else`
`41`	`42`	`{`
`42`	`43`	`char* name = PyString_AsString(self->name);`
`43`	`44`	`if (name == 0)`
`44`	`45`	`return 0;`
`45`	`46`
`46`		`- return PyString_FromFormat("%s.%s", self_->ob_type->tp_name, name);`
	`47`	`+ return PyString_FromFormat("%s.%s.%s", mod, self_->ob_type->tp_name, name);`
`47`	`48`	`}`
`48`	`49`	`}`
`49`	`50`
`@@ -141,10 +142,9 @@ namespace`
`141`	`142`
`142`	`143`	`object module_name = module_prefix();`
`143`	`144`	`if (module_name)`
`144`		`- module_name += '.';`
	`145`	`+ d["__module__"] = module_name;`
`145`	`146`
`146`		`- object result = (object(metatype))(`
`147`		`- module_name + name, make_tuple(base), d);`
	`147`	`+ object result = (object(metatype))(name, make_tuple(base), d);`
`148`	`148`
`149`	`149`	`scope().attr(name) = result;`
`150`	`150`