Repository: tobgu/pyrsistent

Issue:
p*_field prevents pickle from working on a PClass
I would never use pickle in production, but as I was trying to write some strawman example storage code for an example app, I discovered that while I could pickle basic PClasses, I can't pickle any that use pmap_field or pvector_field (and probably others that have similar implementations)

e.g.:

```
>>> class Foo(PClass):
...  v = pvector_field(int)
...
>>> Foo()
Foo(v=IntPVector([]))
>>> dumps(Foo())
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
cPickle.PicklingError: Can't pickle <class 'pyrsistent._field_common.IntPVector'>: attribute lookup pyrsistent._field_common.IntPVector failed
```

The same happens for `pmap_field`.

I guess this is because of the way that those functions generate classes at runtime?

@tobgu if you can let me know what needs done to fix this I can try to submit a PR.

Discussion:

Patch:
diff --git a/pyrsistent/_field_common.py b/pyrsistent/_field_common.py
index 6934978..04231c0 100644
--- a/pyrsistent/_field_common.py
+++ b/pyrsistent/_field_common.py
@@ -2,7 +2,8 @@ from collections import Iterable
 import six
 from pyrsistent._checked_types import (
     CheckedType, CheckedPSet, CheckedPMap, CheckedPVector,
-    optional as optional_type, InvariantException, get_type, wrap_invariant)
+    optional as optional_type, InvariantException, get_type, wrap_invariant,
+    _restore_pickle)
 
 
 def set_fields(dct, bases, name):
@@ -121,12 +122,42 @@ class PTypeError(TypeError):
         self.actual_type = actual_type
 
 
-def _sequence_field(checked_class, suffix, item_type, optional, initial):
+SEQ_FIELD_TYPE_SUFFIXES = {
+    CheckedPVector: "PVector",
+    CheckedPSet: "PSet",
+}
+
+# Global dictionary to hold auto-generated field types: used for unpickling
+_seq_field_types = {}
+
+def _restore_seq_field_pickle(checked_class, item_type, data):
+    """Unpickling function for auto-generated PVec/PSet field types."""
+    type_ = _seq_field_types[checked_class, item_type]
+    return _restore_pickle(type_, data)
+
+def _make_seq_field_type(checked_class, item_type):
+    """Create a subclass of the given checked class with the given item type."""
+    type_ = _seq_field_types.get((checked_class, item_type))
+    if type_ is not None:
+        return type_
+
+    class TheType(checked_class):
+        __type__ = item_type
+
+        def __reduce__(self):
+            return (_restore_seq_field_pickle,
+                    (checked_class, item_type, list(self)))
+
+    suffix = SEQ_FIELD_TYPE_SUFFIXES[checked_class]
+    TheType.__name__ = item_type.__name__.capitalize() + suffix
+    _seq_field_types[checked_class, item_type] = TheType
+    return TheType
+
+def _sequence_field(checked_class, item_type, optional, initial):
     """
     Create checked field for either ``PSet`` or ``PVector``.
 
     :param checked_class: ``CheckedPSet`` or ``CheckedPVector``.
-    :param suffix: Suffix for new type name.
     :param item_type: The required type for the items in the set.
     :param optional: If true, ``None`` can be used as a value for
         this field.
@@ -134,9 +165,7 @@ def _sequence_field(checked_class, suffix, item_type, optional, initial):
 
     :return: A ``field`` containing a checked class.
     """
-    class TheType(checked_class):
-        __type__ = item_type
-    TheType.__name__ = item_type.__name__.capitalize() + suffix
+    TheType = _make_seq_field_type(checked_class, item_type)
 
     if optional:
         def factory(argument):
@@ -164,7 +193,7 @@ def pset_field(item_type, optional=False, initial=()):
 
     :return: A ``field`` containing a ``CheckedPSet`` of the given type.
     """
-    return _sequence_field(CheckedPSet, "PSet", item_type, optional,
+    return _sequence_field(CheckedPSet, item_type, optional,
                            initial)
 
 
@@ -180,13 +209,41 @@ def pvector_field(item_type, optional=False, initial=()):
 
     :return: A ``field`` containing a ``CheckedPVector`` of the given type.
     """
-    return _sequence_field(CheckedPVector, "PVector", item_type, optional,
+    return _sequence_field(CheckedPVector, item_type, optional,
                            initial)
 
 
 _valid = lambda item: (True, "")
 
 
+# Global dictionary to hold auto-generated field types: used for unpickling
+_pmap_field_types = {}
+
+def _restore_pmap_field_pickle(key_type, value_type, data):
+    """Unpickling function for auto-generated PMap field types."""
+    type_ = _pmap_field_types[key_type, value_type]
+    return _restore_pickle(type_, data)
+
+def _make_pmap_field_type(key_type, value_type):
+    """Create a subclass of CheckedPMap with the given key and value types."""
+    type_ = _pmap_field_types.get((key_type, value_type))
+    if type_ is not None:
+        return type_
+
+    class TheMap(CheckedPMap):
+        __key_type__ = key_type
+        __value_type__ = value_type
+
+        def __reduce__(self):
+            return (_restore_pmap_field_pickle,
+                    (self.__key_type__, self.__value_type__, dict(self)))
+
+    TheMap.__name__ = (key_type.__name__.capitalize() +
+                       value_type.__name__.capitalize() + "PMap")
+    _pmap_field_types[key_type, value_type] = TheMap
+    return TheMap
+
+
 def pmap_field(key_type, value_type, optional=False, invariant=PFIELD_NO_INVARIANT):
     """
     Create a checked ``PMap`` field.
@@ -199,11 +256,7 @@ def pmap_field(key_type, value_type, optional=False, invariant=PFIELD_NO_INVARIA
 
     :return: A ``field`` containing a ``CheckedPMap``.
     """
-    class TheMap(CheckedPMap):
-        __key_type__ = key_type
-        __value_type__ = value_type
-    TheMap.__name__ = (key_type.__name__.capitalize() +
-                       value_type.__name__.capitalize() + "PMap")
+    TheMap = _make_pmap_field_type(key_type, value_type)
 
     if optional:
         def factory(argument):


Test patch:
diff --git a/tests/class_test.py b/tests/class_test.py
index f7254b4..0caddba 100644
--- a/tests/class_test.py
+++ b/tests/class_test.py
@@ -2,7 +2,9 @@ from collections import Hashable
 import math
 import pickle
 import pytest
-from pyrsistent import field, InvariantException, PClass, optional, CheckedPVector
+from pyrsistent import (
+    field, InvariantException, PClass, optional, CheckedPVector,
+    pmap_field, pset_field, pvector_field)
 
 
 class Point(PClass):
@@ -11,6 +13,12 @@ class Point(PClass):
     z = field(type=int, initial=0)
 
 
+class TypedContainerObj(PClass):
+    map = pmap_field(str, str)
+    set = pset_field(str)
+    vec = pvector_field(str)
+
+
 def test_evolve_pclass_instance():
     p = Point(x=1, y=2)
     p2 = p.set(x=p.x+2)
@@ -165,6 +173,12 @@ def test_supports_pickling():
     assert isinstance(p2, Point)
 
 
+def test_supports_pickling_with_typed_container_fields():
+    obj = TypedContainerObj(map={'foo': 'bar'}, set=['hello', 'there'], vec=['a', 'b'])
+    obj2 = pickle.loads(pickle.dumps(obj))
+    assert obj == obj2
+
+
 def test_can_remove_optional_member():
     p1 = Point(x=1, y=2)
     p2 = p1.remove('y')
@@ -250,4 +264,4 @@ def test_multiple_global_invariants():
         MultiInvariantGlobal(one=1)
         assert False
     except InvariantException as e:
-        assert e.invariant_errors == (('x', 'y'),)
\ No newline at end of file
+        assert e.invariant_errors == (('x', 'y'),)
diff --git a/tests/record_test.py b/tests/record_test.py
index 9146bd0..b2439e5 100644
--- a/tests/record_test.py
+++ b/tests/record_test.py
@@ -13,6 +13,12 @@ class ARecord(PRecord):
     y = field()
 
 
+class RecordContainingContainers(PRecord):
+    map = pmap_field(str, str)
+    vec = pvector_field(str)
+    set = pset_field(str)
+
+
 def test_create():
     r = ARecord(x=1, y='foo')
     assert r.x == 1
@@ -223,6 +229,11 @@ def test_pickling():
     assert x == y
     assert isinstance(y, ARecord)
 
+def test_supports_pickling_with_typed_container_fields():
+    obj = RecordContainingContainers(
+        map={'foo': 'bar'}, set=['hello', 'there'], vec=['a', 'b'])
+    obj2 = pickle.loads(pickle.dumps(obj))
+    assert obj == obj2
 
 def test_all_invariant_errors_reported():
     class BRecord(PRecord):