Make OrderedSet pickable in more cases, by doing manual weakref.

theano/misc/ordered_set.py

@@ -41,21 +41,16 @@ if MutableSet is not None:
     # SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
     ## {{{ http://code.activestate.com/recipes/576696/ (r5)
     import collections
-    from weakref import proxy
     import weakref
 
     class Link(object):
         __slots__ = 'prev', 'next', 'key', '__weakref__'
 
         def __getstate__(self):
-            # The proxy aren't pickled by default.
-            # We want to restore them.
-            # The way we restore the proxy we create more
-            # proxy, but we do not care.
-            ret = [self.prev, self.next,
-                   isinstance(self.prev, weakref.ProxyType),
-                   isinstance(self.next, weakref.ProxyType),
-                   ]
+            # weakref.proxy don't pickle well, so we use weakref.ref
+            # manually and don't pickle the weakref.
+            # We restore the weakref when we unpickle.
+            ret = [self.prev(), self.next()]
             try:
                 ret.append(self.key)
             except AttributeError:
@@ -63,16 +58,10 @@ if MutableSet is not None:
             return ret
 
         def __setstate__(self, state):
-            if state[2]:
-                self.prev = proxy(state[0])
-            else:
-                self.prev = state[0]
-            if state[3]:
-                self.next = proxy(state[1])
-            else:
-                self.next = state[1]
-            if len(state) == 5:
-                self.key = state[4]
+            self.prev = weakref.ref(state[0])
+            self.next = weakref.ref(state[1])
+            if len(state) == 3:
+                self.key = state[2]
 
     class OrderedSet(collections.MutableSet):
         'Set the remembers the order elements were added'
@@ -94,7 +83,7 @@ if MutableSet is not None:
             # Checks added by IG
             check_deterministic(iterable)
             self.__root = root = Link()         # sentinel node for doubly linked list
-            root.prev = root.next = root
+            root.prev = root.next = weakref.ref(root)
             self.__map = {}                     # key --> link
             if iterable is not None:
                 self |= iterable
@@ -111,8 +100,8 @@ if MutableSet is not None:
                 self.__map[key] = link = Link()
                 root = self.__root
                 last = root.prev
-                link.prev, link.next, link.key = last, root, key
-                last.next = root.prev = proxy(link)
+                link.prev, link.next, link.key = last, weakref.ref(root), key
+                last().next = root.prev = weakref.ref(link)
 
         def union(self, s):
             check_deterministic(s)
@@ -148,24 +137,24 @@ if MutableSet is not None:
             # then removed by updating the links in the predecessor and successors.
             if key in self.__map:
                 link = self.__map.pop(key)
-                link.prev.next = link.next
-                link.next.prev = link.prev
+                link.prev().next = link.next
+                link.next().prev = link.prev
 
         def __iter__(self):
             # Traverse the linked list in order.
             root = self.__root
-            curr = root.next
+            curr = root.next()
             while curr is not root:
                 yield curr.key
-                curr = curr.next
+                curr = curr.next()
 
         def __reversed__(self):
             # Traverse the linked list in reverse order.
             root = self.__root
-            curr = root.prev
+            curr = root.prev()
             while curr is not root:
                 yield curr.key
-                curr = curr.prev
+                curr = curr.prev()
 
         def pop(self, last=True):
             if not self: