ENH: Use np.void for labelarray storage.

This disables most broken ufuncs

tests/test_labelarray.py

@@ -21,6 +21,11 @@ def rotN(l, N):
     return l[N:] + l[:N]
 
 
+def all_ufuncs():
+    ufunc_type = type(np.isnan)
+    return (f for f in vars(np).values() if isinstance(f, ufunc_type))
+
+
 class LabelArrayTestCase(ZiplineTestCase):
 
     @classmethod
@@ -136,7 +141,7 @@ class LabelArrayTestCase(ZiplineTestCase):
         for idx, value in enumerate(arr1d.categories):
             check_arrays(
                 self.strs == value,
-                arr1d.view(type=np.ndarray) == idx,
+                arr1d.as_int_array() == idx,
             )
 
         for shape in (9, 3), (3, 9), (3, 3, 3):
@@ -149,3 +154,30 @@ class LabelArrayTestCase(ZiplineTestCase):
 
             for idx, value in enumerate(arr2d.categories):
                 check_arrays(strs2d == value, codes2d == idx)
+
+    def test_reject_ufuncs(self):
+        """
+        The internal values of a LabelArray should be opaque to numpy ufuncs.
+        """
+        def assert_ufunc_failure(exc):
+            self.assertEqual(str(exc), 'Not implemented for this type')
+
+        l = LabelArray(self.strs, '')
+        ints = np.arange(len(l))
+
+        for func in all_ufuncs():
+            # Different ufuncs vary between returning NotImplemented and
+            # raising a TypeError when provided with unknown dtypes.
+            # This is a bit unfortunate, but still better than silently
+            # accepting an int array.
+            try:
+                if func.nin == 1:
+                    ret = func(l)
+                elif func.nin == 2:
+                    ret = func(l, ints)
+                else:
+                    self.fail("Who added a ternary ufunc !?!")
+            except TypeError as e:
+                assert_ufunc_failure(e)
+            else:
+                self.assertIs(ret, NotImplemented)

zipline/lib/labelarray.py

@@ -19,7 +19,7 @@ from zipline.utils.input_validation import (
     expect_types,
     optional,
 )
-from zipline.utils.numpy_utils import is_object, int64_dtype
+from zipline.utils.numpy_utils import int_dtype_with_size_in_bytes, is_object
 
 from ._factorize import (
     factorize_strings,
@@ -88,6 +88,16 @@ class LabelArray(ndarray):
         supplied, they are left in the order provided.  If sort is False and
         categories is None, categories will be constructed in a random order.
 
+    Attributes
+    ----------
+    categories : ndarray[str]
+        An array containing the unique labels of self.
+    reverse_categories : dict[str -> int]
+        Reverse lookup table for ``categories``. Stores the index in
+        ``categories`` at which each entry each unique entry is found.
+    missing_value : str
+        A sentinel missing value with NaN semantics for comparisons.
+
     Notes
     -----
     Consumers should be cautious when passing instances of LabelArray to numpy
@@ -144,7 +154,26 @@ class LabelArray(ndarray):
             )
         categories.setflags(write=False)
 
-        ret = codes.reshape(values.shape).view(type=cls)
+        return cls._from_codes_and_metadata(
+            codes=codes.reshape(values.shape),
+            categories=categories,
+            reverse_categories=reverse_categories,
+            missing_value=missing_value,
+        )
+
+    @classmethod
+    def _from_codes_and_metadata(cls,
+                                 codes,
+                                 categories,
+                                 reverse_categories,
+                                 missing_value):
+        """
+        View codes as a LabelArray and set LabelArray metadata on the result.
+        """
+        ret = codes.view(
+            type=cls,
+            dtype=np.void(codes.dtype.itemsize),
+        )
         ret._categories = categories
         ret._reverse_categories = reverse_categories
         ret._missing_value = missing_value
@@ -203,24 +232,16 @@ class LabelArray(ndarray):
         self._reverse_categories = getattr(obj, 'reverse_categories', None)
         self._missing_value = getattr(obj, 'missing_value', None)
 
-    def __array_wrap__(self, obj, context=None):
-        """
-        Called by numpy after completion of a ufunc.
-
-        We coerce back into a vanilla ndarray if our dtype changed, since that
-        indicates that our categories are no longer meaningful.
-        """
-        if obj.dtype != self.dtype:
-            return obj.view(type=np.ndarray)
-        return obj
-
     def as_int_array(self):
         """
         Convert self into a regular ndarray of ints.
 
         This is an O(1) operation. It does not copy the underlying data.
         """
-        return self.view(type=ndarray)
+        return self.view(
+            type=ndarray,
+            dtype=int_dtype_with_size_in_bytes(self.itemsize),
+        )
 
     def as_string_array(self):
         """
@@ -228,7 +249,7 @@ class LabelArray(ndarray):
 
         This is an O(N) operation.
         """
-        return self.categories[self]
+        return self.categories[self.as_int_array()]
 
     def as_categorical(self, name=None):
         """
@@ -284,8 +305,7 @@ class LabelArray(ndarray):
             value_code = self.reverse_categories.get(value, None)
             if value_code is None:
                 raise ValueError("%r is not in LabelArray categories." % value)
-            return super(LabelArray, self).__setitem__(indexer, value_code)
-
+            self.as_int_array()[indexer] = value_code
         else:
             raise NotImplementedError(
                 "Setting into a LabelArray with a value of "
@@ -407,10 +427,6 @@ class LabelArray(ndarray):
         # This happens if you call a ufunc on a LabelArray that changes the
         # dtype.  This is generally an indicator that the array has been used
         # incorrectly, and it means we're no longer valid for anything.
-        if self.dtype != int64_dtype:
-            return "Invalid LabelArray: dtype={}, shape={}".format(
-                self.dtype, self.shape
-            )
         repr_lines = repr(self.as_string_array()).splitlines()
         repr_lines[0] = repr_lines[0].replace('array(', 'LabelArray(', 1)
         repr_lines[-1] = repr_lines[-1].rsplit(',', 1)[0] + ')'
@@ -423,27 +439,25 @@ class LabelArray(ndarray):
         Make an empty LabelArray with the same categories as ``self``, filled
         with ``self.missing_value``.
         """
-        out = np.full(
-            shape,
-            self.reverse_categories[self.missing_value],
-            dtype=self.dtype
-        ).view(
-            type=type(self)
+        return type(self)._from_codes_and_metadata(
+            codes=np.full(
+                shape,
+                self.reverse_categories[self.missing_value],
+                dtype=int_dtype_with_size_in_bytes(self.itemsize),
+            ),
+            categories=self.categories,
+            reverse_categories=self.reverse_categories,
+            missing_value=self.missing_value,
         )
 
-        out._categories = self.categories
-        out._reverse_categories = self.reverse_categories
-        out._missing_value = self.missing_value
-
-        return out
-
     def apply(self, f, dtype):
         """
         Map a function elementwise over entries in ``self``.
 
         ``f`` will be applied exactly once to each unique value in ``self``.
         """
-        return np.vectorize(f, otypes=[dtype])(self.categories)[self]
+        vf = np.vectorize(f, otypes=[dtype])
+        return vf(self.categories)[self.as_int_array()]
 
     def startswith(self, prefix):
         """

zipline/utils/numpy_utils.py

@@ -55,6 +55,20 @@ _FILLVALUE_DEFAULTS = {
     datetime64ns_dtype: NaTns,
 }
 
+INT_DTYPES_BY_SIZE_BYTES = {
+    1: dtype('int8'),
+    2: dtype('int16'),
+    4: dtype('int32'),
+    8: dtype('int64'),
+}
+
+
+def int_dtype_with_size_in_bytes(size):
+    try:
+        return INT_DTYPES_BY_SIZE_BYTES[size]
+    except KeyError:
+        raise ValueError("No integral dtype whose size is %d bytes." % size)
+
 
 class NoDefaultMissingValue(Exception):
     pass