Add support for pickling proxy objects. The pickled object will be the object be the object that the proxy points to (__wrapped__). Closes ionelmc#7.

ionelmc · ionelmc · commit 7dca9a29bf04 · 2015-07-02T21:07:17.000+03:00
diff --git a/src/lazy_object_proxy/__init__.py b/src/lazy_object_proxy/__init__.py
@@ -2,7 +2,19 @@
 
 __all__ = "Proxy",
 
+try:
+    import copy_reg as copyreg
+except ImportError:
+    import copyreg
+
+from .utils import identity
+
+copyreg.constructor(identity)
+
 try:
     from .cext import Proxy
+    from .cext import identity
 except ImportError:
     from .slots import Proxy
+else:
+    copyreg.constructor(identity)
diff --git a/src/lazy_object_proxy/cext.c b/src/lazy_object_proxy/cext.c
@@ -33,6 +33,26 @@ typedef struct {
 
 PyTypeObject Proxy_Type;
 
+
+/* ------------------------------------------------------------------------- */
+
+static PyObject *identity_ref = NULL;
+static PyObject *
+identity(PyObject *self, PyObject *value)
+{
+    Py_INCREF(value);
+    return value;
+}
+
+/* ------------------------------------------------------------------------- */
+
+PyDoc_STRVAR(identity_doc, "Indentity function: returns the single argument.");
+
+static struct PyMethodDef module_functions[] = {
+    {"identity", identity, METH_O, identity_doc},
+    {NULL,       NULL}
+};
+
 /* ------------------------------------------------------------------------- */
 
 static PyObject *Proxy__ensure_wrapped(ProxyObject *self)
@@ -844,6 +864,15 @@ static PyObject *Proxy_reversed(
             self->wrapped, NULL);
 }
 
+/* ------------------------------------------------------------------------- */
+static PyObject *Proxy_reduce(
+        ProxyObject *self, PyObject *args)
+{
+    Proxy__ENSURE_WRAPPED_OR_RETURN_NULL(self);
+
+    return Py_BuildValue("(O(O))", identity_ref, self->wrapped);
+}
+
 /* ------------------------------------------------------------------------- */
 
 #if PY_MAJOR_VERSION >= 3
@@ -1227,6 +1256,8 @@ static PyMethodDef Proxy_methods[] = {
                     METH_VARARGS , 0 },
     { "__bytes__",  (PyCFunction)Proxy_bytes, METH_NOARGS, 0 },
     { "__reversed__", (PyCFunction)Proxy_reversed, METH_NOARGS, 0 },
+    { "__reduce__", (PyCFunction)Proxy_reduce, METH_NOARGS, 0 },
+    { "__reduce_ex__", (PyCFunction)Proxy_reduce, METH_O, 0 },
 #if PY_MAJOR_VERSION >= 3
     { "__round__",  (PyCFunction)Proxy_round, METH_NOARGS, 0 },
 #endif
@@ -1308,26 +1339,27 @@ PyTypeObject Proxy_Type = {
 #if PY_MAJOR_VERSION >= 3
 static struct PyModuleDef moduledef = {
     PyModuleDef_HEAD_INIT,
-    "cext",              /* m_name */
-    NULL,                /* m_doc */
-    -1,                  /* m_size */
-    NULL,                /* m_methods */
-    NULL,                /* m_reload */
-    NULL,                /* m_traverse */
-    NULL,                /* m_clear */
-    NULL,                /* m_free */
+    "lazy_object_proxy.cext", /* m_name */
+    NULL,                     /* m_doc */
+    -1,                       /* m_size */
+    module_functions,         /* m_methods */
+    NULL,                     /* m_reload */
+    NULL,                     /* m_traverse */
+    NULL,                     /* m_clear */
+    NULL,                     /* m_free */
 };
 #endif
 
 static PyObject *
 moduleinit(void)
 {
     PyObject *module;
+    PyObject *dict;
 
 #if PY_MAJOR_VERSION >= 3
     module = PyModule_Create(&moduledef);
 #else
-    module = Py_InitModule3("cext", NULL, NULL);
+    module = Py_InitModule3("lazy_object_proxy.cext", module_functions, NULL);
 #endif
 
     if (module == NULL)
@@ -1339,6 +1371,14 @@ moduleinit(void)
     if (PyType_Ready(&Proxy_Type) < 0)
         return NULL;
 
+    dict = PyModule_GetDict(module);
+    if (dict == NULL)
+        return NULL;
+    identity_ref = PyDict_GetItemString(dict, "identity");
+    if (identity_ref == NULL)
+        return NULL;
+    Py_INCREF(identity_ref);
+
     Py_INCREF(&Proxy_Type);
     PyModule_AddObject(module, "Proxy",
             (PyObject *)&Proxy_Type);
diff --git a/src/lazy_object_proxy/simple.py b/src/lazy_object_proxy/simple.py
@@ -3,22 +3,14 @@
 from .compat import PY2
 from .compat import PY3
 from .compat import with_metaclass
-
-
-class cached_property(object):
-    def __init__(self, func):
-        self.func = func
-
-    def __get__(self, obj, cls):
-        if obj is None:
-            return self
-        value = obj.__dict__[self.func.__name__] = self.func(obj)
-        return value
+from .utils import cached_property
+from .utils import identity
 
 
 def make_proxy_method(code):
     def proxy_wrapper(self, *args):
         return code(self.__wrapped__, *args)
+
     return proxy_wrapper
 
 
@@ -86,7 +78,6 @@ def __wrapped__(self):
         else:
             raise ValueError("Proxy hasn't been initiated: __factory__ is missing.")
 
-
     __name__ = property(make_proxy_method(operator.attrgetter('__name__')))
     __class__ = property(make_proxy_method(operator.attrgetter('__class__')))
     __annotations__ = property(make_proxy_method(operator.attrgetter('__anotations__')))
@@ -240,3 +231,9 @@ def __exit__(self, *args, **kwargs):
 
     def __call__(self, *args, **kwargs):
         return self.__wrapped__(*args, **kwargs)
+
+    def __reduce__(self):
+        return identity, (self.__wrapped__,)
+
+    def __reduce_ex__(self, protocol):
+        return identity, (self.__wrapped__,)
diff --git a/src/lazy_object_proxy/slots.py b/src/lazy_object_proxy/slots.py
@@ -3,6 +3,7 @@
 from .compat import PY2
 from .compat import PY3
 from .compat import with_metaclass
+from .utils import identity
 
 
 class _ProxyMethods(object):
@@ -81,7 +82,8 @@ def __init__(self, factory):
         object.__setattr__(self, '__factory__', factory)
 
     @property
-    def __wrapped__(self, __getattr__=object.__getattribute__, __setattr__=object.__setattr__, __delattr__=object.__delattr__):
+    def __wrapped__(self, __getattr__=object.__getattribute__, __setattr__=object.__setattr__,
+                    __delattr__=object.__delattr__):
         try:
             return __getattr__(self, '__target__')
         except AttributeError:
@@ -395,3 +397,9 @@ def __iter__(self):
 
     def __call__(self, *args, **kwargs):
         return self.__wrapped__(*args, **kwargs)
+
+    def __reduce__(self):
+        return identity, (self.__wrapped__,)
+
+    def __reduce_ex__(self, protocol):
+        return identity, (self.__wrapped__,)
diff --git a/src/lazy_object_proxy/utils.py b/src/lazy_object_proxy/utils.py
@@ -0,0 +1,13 @@
+def identity(obj):
+    return obj
+
+
+class cached_property(object):
+    def __init__(self, func):
+        self.func = func
+
+    def __get__(self, obj, cls):
+        if obj is None:
+            return self
+        value = obj.__dict__[self.func.__name__] = self.func(obj)
+        return value
diff --git a/tests/test_lazy_object_proxy.py b/tests/test_lazy_object_proxy.py
