COMPAT: ensure proper extension dtype's don't pickle the cache

Author: jreback

pandas/core/dtypes/dtypes.py

@@ -24,6 +24,7 @@ class ExtensionDtype(object):
     isbuiltin = 0
     isnative = 0
     _metadata = []
+    _cache = {}
 
     def __unicode__(self):
         return self.name
@@ -71,6 +72,15 @@ def __eq__(self, other):
     def __ne__(self, other):
         return not self.__eq__(other)
 
+    def __getstate__(self):
+        # pickle support; we don't want to pickle the cache
+        return {k: getattr(self, k, None) for k in self._metadata}
+
+    @classmethod
+    def reset_cache(cls):
+        """ clear the cache """
+        cls._cache = {}
+
     @classmethod
     def is_dtype(cls, dtype):
         """ Return a boolean if the passed type is an actual dtype that
@@ -111,6 +121,7 @@ class CategoricalDtype(ExtensionDtype):
     str = '|O08'
     base = np.dtype('O')
     _cache = {}
+    _attributes = ['name']
 
     def __new__(cls):
 
@@ -423,6 +434,12 @@ def __new__(cls, subtype=None):
             except TypeError:
                 raise ValueError("could not construct IntervalDtype")
 
+        if subtype is None:
+            # pickle compat
+            u = object.__new__(cls)
+            u.subtype = None
+            return u
+
         try:
             return cls._cache[str(subtype)]
         except KeyError:

pandas/tests/dtypes/test_dtypes.py

@@ -23,6 +23,9 @@
 
 class Base(object):
 
+    def setUp(self):
+        self.dtype = self.create()
+
     def test_hash(self):
         hash(self.dtype)
 
@@ -37,14 +40,38 @@ def test_numpy_informed(self):
         assert not np.str_ == self.dtype
 
     def test_pickle(self):
+        # make sure our cache is NOT pickled
+
+        # clear the cache
+        type(self.dtype).reset_cache()
+        assert not len(self.dtype._cache)
+
+        # force back to the cache
         result = tm.round_trip_pickle(self.dtype)
+        assert not len(self.dtype._cache)
         assert result == self.dtype
 
 
 class TestCategoricalDtype(Base, tm.TestCase):
 
-    def setUp(self):
-        self.dtype = CategoricalDtype()
+    def create(self):
+        return CategoricalDtype()
+
+    def test_pickle(self):
+        # make sure our cache is NOT pickled
+
+        # clear the cache
+        type(self.dtype).reset_cache()
+        assert not len(self.dtype._cache)
+
+        # force back to the cache
+        result = tm.round_trip_pickle(self.dtype)
+
+        # we are a singular object so we are added
+        # back to the cache upon unpickling
+        # this is to ensure object identity
+        assert len(self.dtype._cache) == 1
+        assert result == self.dtype
 
     def test_hash_vs_equality(self):
         # make sure that we satisfy is semantics
@@ -95,8 +122,8 @@ def test_basic(self):
 
 class TestDatetimeTZDtype(Base, tm.TestCase):
 
-    def setUp(self):
-        self.dtype = DatetimeTZDtype('ns', 'US/Eastern')
+    def create(self):
+        return DatetimeTZDtype('ns', 'US/Eastern')
 
     def test_hash_vs_equality(self):
         # make sure that we satisfy is semantics
@@ -211,8 +238,8 @@ def test_empty(self):
 
 class TestPeriodDtype(Base, tm.TestCase):
 
-    def setUp(self):
-        self.dtype = PeriodDtype('D')
+    def create(self):
+        return PeriodDtype('D')
 
     def test_construction(self):
         with pytest.raises(ValueError):
@@ -340,9 +367,22 @@ def test_not_string(self):
 
 class TestIntervalDtype(Base, tm.TestCase):
 
-    # TODO: placeholder
-    def setUp(self):
-        self.dtype = IntervalDtype('int64')
+    def create(self):
+        return IntervalDtype('int64')
+
+    def test_hash_vs_equality(self):
+        # make sure that we satisfy is semantics
+        dtype = self.dtype
+        dtype2 = IntervalDtype('int64')
+        dtype3 = IntervalDtype(dtype2)
+        assert dtype == dtype2
+        assert dtype2 == dtype
+        assert dtype3 == dtype
+        assert dtype is dtype2
+        assert dtype2 is dtype
+        assert dtype3 is dtype
+        assert hash(dtype) == hash(dtype2)
+        assert hash(dtype) == hash(dtype3)
 
     def test_construction(self):
         with pytest.raises(ValueError):
@@ -445,3 +485,14 @@ def test_basic_dtype(self):
         assert not is_interval_dtype(np.object_)
         assert not is_interval_dtype(np.int64)
         assert not is_interval_dtype(np.float64)
+
+    def test_caching(self):
+        IntervalDtype.reset_cache()
+        dtype = IntervalDtype("int64")
+        assert len(IntervalDtype._cache) == 1
+
+        IntervalDtype("interval")
+        assert len(IntervalDtype._cache) == 1
+
+        tm.round_trip_pickle(dtype)
+        assert len(IntervalDtype._cache) == 1