refactor dtypes

Author: jreback

pandas/core/arrays/mask/_base.py

@@ -7,8 +7,42 @@
 
 import copy
 import numpy as np
+from pandas import compat
 from pandas.core.missing import isna
 from pandas.core.arrays.base import ExtensionArray
+from pandas.api.extensions import ExtensionDtype
+
+
+class BoolDtype(ExtensionDtype):
+
+    type = np.bool_
+    kind = 'b'
+    name = 'bool'
+
+    @classmethod
+    def construct_from_string(cls, string):
+        if string == cls.name:
+            return cls()
+        else:
+            raise TypeError("Cannot construct a '{}' from "
+                            "'{}'".format(cls, string))
+
+    def _is_boolean(self):
+        return True
+
+    def __hash__(self):
+        return hash(str(self))
+
+    def __eq__(self, other):
+        # compare == to np.dtype('bool')
+        if isinstance(other, compat.string_types):
+            return other == self.name
+        elif other is self:
+            return True
+        elif isinstance(other, np.dtype):
+            return other == 'bool'
+        else:
+            return hash(self) == hash(other)
 
 
 class BoolArray(ExtensionArray):

pandas/core/arrays/mask/_numpy.py

@@ -4,34 +4,19 @@
 
 import numpy as np
 
-from pandas.core.arrays.mask._base import BoolArray
-from pandas.api.extensions import (
-    ExtensionDtype, register_extension_dtype)
+from pandas.core.arrays.mask._base import BoolDtype, BoolArray
+from pandas.api.extensions import register_extension_dtype, take
 
 
 @register_extension_dtype
-class NumpyBoolDtype(ExtensionDtype):
+class NumpyBoolDtype(BoolDtype):
 
-    type = np.bool_
-    kind = 'b'
-    name = 'bool'
     na_value = np.nan
 
-    @classmethod
-    def construct_from_string(cls, string):
-        if string == cls.name:
-            return cls()
-        else:
-            raise TypeError("Cannot construct a '{}' from "
-                            "'{}'".format(cls, string))
-
     @classmethod
     def construct_array_type(cls):
         return NumpyBoolArray
 
-    def _is_boolean(self):
-        return True
-
 
 class NumpyBoolArray(BoolArray):
     """Generic class which can be used to represent missing data.
@@ -79,5 +64,13 @@ def reshape(self, shape, **kwargs):
     def astype(self, dtype, copy=False):
         return np.array(self, copy=False).astype(dtype, copy=copy)
 
-    def take(self, indicies, **kwargs):
-        return np.array(self, copy=False).take(indicies)
+    def take(self, indices, allow_fill=False, fill_value=None, axis=None):
+        # TODO: had to add axis here
+        data = self._data
+
+        if allow_fill and fill_value is None:
+            fill_value = self.dtype.na_value
+
+        result = take(data, indices, fill_value=fill_value,
+                      allow_fill=allow_fill)
+        return self._from_sequence(result, dtype=self.dtype)

pandas/core/arrays/mask/_pyarrow.py

@@ -12,8 +12,8 @@
 
 from pandas.api.types import is_scalar
 from pandas.api.extensions import (
-    ExtensionDtype, register_extension_dtype, take)
-from pandas.core.arrays.mask._base import BoolArray
+    register_extension_dtype, take)
+from pandas.core.arrays.mask._base import BoolDtype, BoolArray
 
 
 # we require pyarrow >= 0.10.0
@@ -27,28 +27,14 @@
 
 
 @register_extension_dtype
-class ArrowBoolDtype(ExtensionDtype):
+class ArrowBoolDtype(BoolDtype):
 
-    type = np.bool_
-    kind = 'b'
-    name = 'arrow_bool'
     na_value = pa.NULL
 
-    @classmethod
-    def construct_from_string(cls, string):
-        if string == cls.name:
-            return cls()
-        else:
-            raise TypeError("Cannot construct a '{}' from "
-                            "'{}'".format(cls, string))
-
     @classmethod
     def construct_array_type(cls):
         return ArrowBoolArray
 
-    def _is_boolean(self):
-        return True
-
 
 class ArrowBoolArray(BoolArray):