fix pickling of multidicts (#77)

* fix pickling of multidicts

* add unit test

* fix unittest and simplify code

* disable pickling of proxies as we can't unpickle them correctly

* flake

* changes based on review

use py3.6 exception type

* fix test

* fix unittest

* remove old import

* pep fix

Author: thehesiod

multidict/_multidict.pyx

@@ -3,7 +3,6 @@ from __future__ import absolute_import
 import sys
 from collections import abc
 from collections.abc import Iterable, Set
-from operator import itemgetter
 
 from cpython.object cimport PyObject_Str
 
@@ -216,37 +215,34 @@ cdef class _Base:
 
 
 cdef class MultiDictProxy(_Base):
+    _proxy_classes = (MultiDict, MultiDictProxy)
+    _base_class = MultiDict
 
     def __init__(self, arg):
         cdef _Base base
-        if not isinstance(arg, (MultiDict, MultiDictProxy)):
+        if not isinstance(arg, self._proxy_classes):
             raise TypeError(
-                'ctor requires MultiDict or MultiDictProxy instance'
+                'ctor requires {} instance'
                 ', not {}'.format(
+                    ' or '.join(self._proxy_classes),
                     type(arg)))
 
         base = arg
         self._impl = base._impl
 
+    def __reduce__(self):
+        raise TypeError("can't pickle {} objects".format(self.__class__.__name__))
+
     def copy(self):
         """Return a copy of itself."""
-        return MultiDict(self)
+        return self._base_class(self)
 
 abc.Mapping.register(MultiDictProxy)
 
 
 cdef class CIMultiDictProxy(MultiDictProxy):
-
-    def __init__(self, arg):
-        cdef _Base base
-        if not isinstance(arg, (CIMultiDict, CIMultiDictProxy)):
-            raise TypeError(
-                'ctor requires CIMultiDict or CIMultiDictProxy instance'
-                ', not {}'.format(
-                    type(arg)))
-
-        base = arg
-        self._impl = base._impl
+    _proxy_classes = (CIMultiDict, CIMultiDictProxy)
+    _base_class = CIMultiDict
 
     cdef str _title(self, s):
         typ = type(s)
@@ -256,10 +252,6 @@ cdef class CIMultiDictProxy(MultiDictProxy):
             return PyObject_Str(s)
         return s.title()
 
-    def copy(self):
-        """Return a copy of itself."""
-        return CIMultiDict(self)
-
 
 abc.Mapping.register(CIMultiDictProxy)
 
@@ -282,9 +274,14 @@ cdef class MultiDict(_Base):
 
     def __init__(self, *args, **kwargs):
         self._impl = _Impl()
-
         self._extend(args, kwargs, 'MultiDict', True)
 
+    def __reduce__(self):
+        return (
+            self.__class__,
+            tuple(self.items()),
+        )
+
     cdef _extend(self, tuple args, dict kwargs, name, bint do_add):
         cdef _Pair item
         cdef object key

multidict/_multidict_py.py

@@ -155,6 +155,10 @@ def __init__(self, arg):
 
         self._impl = arg._impl
 
+    def __reduce__(self):
+        raise TypeError("can't pickle {} objects".format(
+            self.__class__.__name__))
+
     def copy(self):
         """Return a copy of itself."""
         return MultiDict(self.items())

tests/test_multidict.py

@@ -1,5 +1,6 @@
 import sys
 import unittest
+import pickle
 
 import multidict
 from multidict._multidict import (MultiDictProxy,
@@ -327,6 +328,20 @@ def test_values__repr__(self):
         self.assertEqual(repr(d.values()),
                          "_ValuesView('value1', 'value2')")
 
+    def test_pickle(self):
+        d = self.make_dict([('a', 1), ('a', 2)])
+
+        if isinstance(d, (MultiDictProxy, _MultiDictProxy)):
+            with self.assertRaises(TypeError):
+                pbytes = pickle.dumps(d)
+
+            return
+        else:
+            pbytes = pickle.dumps(d)
+
+        obj = pickle.loads(pbytes)
+        self.assertEqual(dict(d), dict(obj))
+
 
 class _CIMultiDictTests(_Root):