BUG: itertools objects are actually picklable

Author: Joe Jevnik

cloudpickle/cloudpickle.py

@@ -284,16 +284,6 @@ def save_buffer(self, obj):
 
         dispatch[buffer] = save_buffer
 
-    def save_unsupported(self, obj):
-        raise pickle.PicklingError("Cannot pickle objects of type %s" % type(obj))
-
-    dispatch[types.GeneratorType] = save_unsupported
-
-    # itertools objects do not pickle!
-    for v in itertools.__dict__.values():
-        if type(v) is type:
-            dispatch[v] = save_unsupported
-
     def save_module(self, obj):
         """
         Save a module as an import

tests/cloudpickle_test.py

@@ -347,15 +347,6 @@ def test_ufunc(self):
         else:  # skip if scipy is not available
             pass
 
-    def test_save_unsupported(self):
-        sio = StringIO()
-        pickler = cloudpickle.CloudPickler(sio, 2)
-
-        with pytest.raises(pickle.PicklingError) as excinfo:
-            pickler.save_unsupported("test")
-
-        assert "Cannot pickle objects of type" in str(excinfo.value)
-
     def test_loads_namespace(self):
         obj = 1, 2, 3, 4
         returned_obj = cloudpickle.loads(cloudpickle.dumps(obj))
@@ -883,6 +874,20 @@ def test_unhashable_function(self):
         self.assertEquals(depickled_method('a'), 1)
         self.assertEquals(depickled_method('b'), None)
 
+    def test_itertools_count(self):
+        counter = itertools.count(1, step=2)
+
+        # advance the counter a bit
+        next(counter)
+        next(counter)
+
+        new_counter = pickle_depickle(counter, protocol=self.protocol)
+
+        self.assertTrue(counter is not new_counter)
+
+        for _ in range(10):
+            self.assertEqual(next(counter), next(new_counter))
+
 
 class Protocol2CloudPickleTest(CloudPickleTest):