From 620d7648b0b22e4b2509506327a42e95358cf599 Mon Sep 17 00:00:00 2001
From: Scott Sanderson <ssanderson@quantopian.com>
Date: Tue, 3 May 2016 19:53:52 -0400
Subject: [PATCH] BUG: Tests/bugfixes for LabelArray slicing.

- Fixes a bug where __setitem__ was not called when setting with a slice
  on Python 2 (__setslice__ was called instead), which caused strange
  behavior when setting an empty string.  This is fixed by overriding
  __setslice__ and forwarding to __setitem__.

- Fixes a bug where __getitem__ returned an instance of np.void when
  returning a scalar.  We now correctly return an entry from our
  categoricals.
---
 tests/test_labelarray.py  | 75 +++++++++++++++++++++++++++++++++++++++
 zipline/lib/labelarray.py | 27 +++++++++++---
 2 files changed, 98 insertions(+), 4 deletions(-)

diff --git a/tests/test_labelarray.py b/tests/test_labelarray.py
index 9c2462e6..a5a41157 100644
--- a/tests/test_labelarray.py
+++ b/tests/test_labelarray.py
@@ -222,6 +222,8 @@ class LabelArrayTestCase(ZiplineTestCase):
     def test_reject_ufuncs(self):
         """
         The internal values of a LabelArray should be opaque to numpy ufuncs.
+
+        Test that all unfuncs fail.
         """
         def assert_ufunc_failure(exc):
             self.assertEqual(str(exc), 'Not implemented for this type')
@@ -245,3 +247,76 @@ class LabelArrayTestCase(ZiplineTestCase):
                 assert_ufunc_failure(e)
             else:
                 self.assertIs(ret, NotImplemented)
+
+    @parameter_space(
+        __fail_fast=True,
+        val=['', 'a', 'not in the array', None],
+        missing_value=['', 'a', 'not in the array', None],
+    )
+    def test_setitem_scalar(self, val, missing_value):
+        arr = LabelArray(self.strs, missing_value=missing_value)
+
+        if not arr.has_label(val):
+            self.assertTrue(
+                (val == 'not in the array')
+                or (val is None and missing_value is not None)
+            )
+            for slicer in [(0, 0), (0, 1), 1]:
+                with self.assertRaises(ValueError):
+                    arr[slicer] = val
+            return
+
+        arr[0, 0] = val
+        self.assertEqual(arr[0, 0], val)
+
+        arr[0, 1] = val
+        self.assertEqual(arr[0, 1], val)
+
+        arr[1] = val
+        if val == missing_value:
+            self.assertTrue(arr.is_missing()[1].all())
+        else:
+            self.assertTrue((arr[1] == val).all())
+            self.assertTrue((arr[1].as_string_array() == val).all())
+
+        arr[:, -1] = val
+        if val == missing_value:
+            self.assertTrue(arr.is_missing()[:, -1].all())
+        else:
+            self.assertTrue((arr[:, -1] == val).all())
+            self.assertTrue((arr[:, -1].as_string_array() == val).all())
+
+        arr[:] = val
+        if val == missing_value:
+            self.assertTrue(arr.is_missing().all())
+        else:
+            self.assertFalse(arr.is_missing().any())
+            self.assertTrue((arr == val).all())
+
+    def test_setitem_array(self):
+        arr = LabelArray(self.strs, missing_value=None)
+        orig_arr = arr.copy()
+
+        # Write a row.
+        self.assertFalse(
+            (arr[0] == arr[1]).all(),
+            "This test doesn't test anything because rows 0"
+            " and 1 are already equal!"
+        )
+        arr[0] = arr[1]
+        for i in range(arr.shape[1]):
+            self.assertEqual(arr[0, i], arr[1, i])
+
+        # Write a column.
+        self.assertFalse(
+            (arr[:, 0] == arr[:, 1]).all(),
+            "This test doesn't test anything because columns 0"
+            " and 1 are already equal!"
+        )
+        arr[:, 0] = arr[:, 1]
+        for i in range(arr.shape[0]):
+            self.assertEqual(arr[i, 0], arr[i, 1])
+
+        # Write the whole array.
+        arr[:] = orig_arr
+        check_arrays(arr, orig_arr)
diff --git a/zipline/lib/labelarray.py b/zipline/lib/labelarray.py
index 1c233e66..03ca46f5 100644
--- a/zipline/lib/labelarray.py
+++ b/zipline/lib/labelarray.py
@@ -189,10 +189,7 @@ class LabelArray(ndarray):
         """
         View codes as a LabelArray and set LabelArray metadata on the result.
         """
-        ret = codes.view(
-            type=cls,
-            dtype=np.void(codes.dtype.itemsize),
-        )
+        ret = codes.view(type=cls, dtype=np.void)
         ret._categories = categories
         ret._reverse_categories = reverse_categories
         ret._missing_value = missing_value
@@ -213,6 +210,9 @@ class LabelArray(ndarray):
         # This is a property because it should be immutable.
         return self._missing_value
 
+    def has_label(self, value):
+        return value in self.reverse_categories
+
     def __array_finalize__(self, obj):
         """
         Called by Numpy after array construction.
@@ -333,6 +333,25 @@ class LabelArray(ndarray):
                 ),
             )
 
+    def __setslice__(self, i, j, sequence):
+        """
+        This method was deprecated in Python 2.0. It predates slice objects,
+        but Python 2.7.11 still uses it if you implement it, which ndarray
+        does.  In newer Pythons, __setitem__ is always called, but we need to
+        manuallly forward in py2.
+        """
+        self.__setitem__(slice(i, j), sequence)
+
+    def __getitem__(self, indexer):
+        result = super(LabelArray, self).__getitem__(indexer)
+        if result.ndim:
+            # Result is still a LabelArray, so we can just return it.
+            return result
+
+        # Result is a scalar value, which will be an instance of np.void.
+        # Map it back to one of our category entries.
+        return self.categories[result.view(int)]
+
     def is_missing(self):
         """
         Like isnan, but checks for locations where we store missing values.