diff --git a/tests/test_labelarray.py b/tests/test_labelarray.py index 2f4a8967..f7a32f40 100644 --- a/tests/test_labelarray.py +++ b/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) diff --git a/zipline/lib/labelarray.py b/zipline/lib/labelarray.py index ffd43c2e..f33a9021 100644 --- a/zipline/lib/labelarray.py +++ b/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): """ diff --git a/zipline/utils/numpy_utils.py b/zipline/utils/numpy_utils.py index a254fff8..63be0aa5 100644 --- a/zipline/utils/numpy_utils.py +++ b/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