First pass at implementation (needs refactor)

Shared methods to simplify implementation

Added try...except for non-ndarray Series _set_with

Minor fixups

Reverted changes; created new module for mask

Reverted test from master

Revert unnecessary changes from master

Added iter to mask

Added more dunders

More properties / methods

Author: WillAyd

pandas/core/arrays/_mask.py

@@ -0,0 +1,103 @@
+import numpy as np
+
+
+class NAMask():
+    """Generic class which can be used to represent missing data.
+
+    Will use bitarray if available; otherwise will use numpy."""
+
+    def __init__(self, mask):
+        """
+        Parameters
+        ----------
+        mask : numpy array
+            Mask of missing values.
+        """
+
+        self._has_bitarray = False
+        try:
+            import bitarray
+            globals()['bitarray'] = bitarray
+            self._has_bitarray = True
+            self._data = self._numpy_to_bitarray(mask)
+        except (ImportError, ModuleNotFoundError):
+            self._data = mask.astype(bool, copy=False)
+
+    def _numpy_to_bitarray(self, arr):
+        bit_arr = bitarray()
+        bit_arr.pack(arr.astype(bool, copy=False))
+
+    def _bitarray_to_numpy(self, arr):
+        return np.fromstring(arr.unpack(), dtype=bool)
+
+    def __getitem__(self, item):
+        if self._has_bitarray:
+            raise NotImplementedError
+
+        return self._data[item]
+
+    def __setitem__(self, key, value):
+        if self._has_bitarray:
+            raise NotImplementedError
+
+        self._data[key] = value
+
+    def __array__(self):
+        if self._has_bitarray:
+            raise NotImplementedError
+
+        return self._data
+
+    def __iter__(self):
+        for i in range(len(self._data)):
+            yield self._data[i]
+
+    def __invert__(self):
+        if self._has_bitarray:
+            raise NotImplementedError
+
+        return type(self)(~self._data)
+
+    def __or__(self, other):
+        if self._has_bitarray:
+            raise NotImplementedError
+
+        return self._data.__or__(other)
+
+    def __ior__(self, other):
+        if self._has_bitarray:
+            raise NotImplementedError
+
+        return self._data | other
+
+    @property
+    def nbytes(self):
+        if self._has_bitarray:
+            return self._data.buffer_info()[1]
+
+        return self._data.nbytes
+
+    @property
+    def size(self):
+        if self._has_bitarray:
+            raise NotImplementedError
+
+        return self._data.size        
+
+    def astype(self, dtype, copy=False):
+        if self._has_bitarray:
+            raise NotImplementedError
+
+        return self._data.astype(dtype, copy=copy)
+
+    def any(self):
+        return self._data.any()
+
+    def copy(self):
+        return type(self)(self._data.copy())
+
+    def sum(self):
+        if self._has_bitarray:
+            raise NotImplementedError
+
+        return self._data.sum()

pandas/core/arrays/integer.py

@@ -19,7 +19,11 @@
 
 from pandas.core import nanops
 from pandas.core.arrays import ExtensionArray, ExtensionOpsMixin
+<<<<<<< HEAD
 from pandas.core.tools.numeric import to_numeric
+=======
+from pandas.core.arrays._mask import NAMask
+>>>>>>> 384287e71... Revert unnecessary changes from master
 
 
 class _IntegerDtype(ExtensionDtype):
@@ -287,7 +291,8 @@ def __init__(self, values, mask, copy=False):
                 and is_integer_dtype(values.dtype)):
             raise TypeError("values should be integer numpy array. Use "
                             "the 'integer_array' function instead")
-        if not (isinstance(mask, np.ndarray) and is_bool_dtype(mask.dtype)):
+        if not (isinstance(mask, NAMask) or (
+                isinstance(mask, np.ndarray) and is_bool_dtype(mask.dtype))):
             raise TypeError("mask should be boolean numpy array. Use "
                             "the 'integer_array' function instead")
 
@@ -296,7 +301,7 @@ def __init__(self, values, mask, copy=False):
             mask = mask.copy()
 
         self._data = values
-        self._mask = mask
+        self._mask = NAMask(mask)
 
     @classmethod
     def _from_sequence(cls, scalars, dtype=None, copy=False):
@@ -323,7 +328,22 @@ def __getitem__(self, item):
             if self._mask[item]:
                 return self.dtype.na_value
             return self._data[item]
+<<<<<<< HEAD
         return type(self)(self._data[item], self._mask[item])
+=======
+
+<<<<<<< HEAD
+        return type(self)(self._data[item],
+<<<<<<< HEAD
+                          mask=_bitarray_to_numpy(self._mask)[item],
+                          dtype=self.dtype)
+>>>>>>> 2ff4b0907... First pass at implementation (needs refactor)
+=======
+                          mask=_bitarray_to_numpy(self._mask)[item])
+>>>>>>> e085674ac... Reverted changes; created new module for mask
+=======
+        return type(self)(self._data[item], self._mask[item])
+>>>>>>> 384287e71... Revert unnecessary changes from master
 
     def _coerce_to_ndarray(self):
         """

pandas/core/series.py

@@ -1067,9 +1067,18 @@ def _set_with(self, key, value):
                 except Exception:
                     pass
 
+<<<<<<< HEAD
+<<<<<<< HEAD
             if is_scalar(key):
                 key = [key]
             elif not isinstance(key, (list, Series, np.ndarray)):
+=======
+            if not isinstance(key, (
+                    list, Series, np.ndarray, Series)):
+>>>>>>> cae87e9df... Added try...except for non-ndarray Series _set_with
+=======
+            if not isinstance(key, (list, Series, np.ndarray, Series)):
+>>>>>>> 65893ae93... Minor fixups
                 try:
                     key = list(key)
                 except Exception:

pandas/util/testing.py

@@ -15,6 +15,15 @@
 import traceback
 import warnings
 
+<<<<<<< HEAD
+=======
+from datetime import datetime
+from functools import wraps
+from contextlib import contextmanager
+
+from bitarray import bitarray
+from numpy.random import randn, rand
+>>>>>>> 2ff4b0907... First pass at implementation (needs refactor)
 import numpy as np
 from numpy.random import rand, randn
 
@@ -1195,15 +1204,33 @@ def assert_extension_array_equal(left, right, check_dtype=True,
     A mask of missing values is computed for each and checked to match.
     The remaining all-valid values are cast to object dtype and checked.
     """
+<<<<<<< HEAD
     assert isinstance(left, ExtensionArray), 'left is not an ExtensionArray'
     assert isinstance(right, ExtensionArray), 'right is not an ExtensionArray'
     if check_dtype:
         assert_attr_equal('dtype', left, right, obj='ExtensionArray')
+<<<<<<< HEAD
+=======
+    assert isinstance(left, ExtensionArray)
+    assert left.dtype == right.dtype
+    left_na = left.isna()
+    right_na = right.isna()
+
+    # TODO - maybe generate dedicated method for bitarray comparison?
+    if isinstance(left_na, bitarray):
+        left_na = np.fromstring(left_na.unpack(), dtype=bool)
+    if isinstance(right_na, bitarray):
+        right_na = np.fromstring(right_na.unpack(), dtype=bool)
+
+    assert_numpy_array_equal(left_na, right_na)
+>>>>>>> 2ff4b0907... First pass at implementation (needs refactor)
 
     if hasattr(left, "asi8") and type(right) == type(left):
         # Avoid slow object-dtype comparisons
         assert_numpy_array_equal(left.asi8, right.asi8)
         return
+=======
+>>>>>>> 36256a718... Reverted test from master
 
     left_na = np.asarray(left.isna())
     right_na = np.asarray(right.isna())