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