MAINT: Add nearest_unequal_elements..

tests/utils/test_pandas_utils.py

@@ -0,0 +1,79 @@
+"""
+Tests for zipline/utils/pandas_utils.py
+"""
+import pandas as pd
+
+from zipline.testing import parameter_space, ZiplineTestCase
+from zipline.utils.pandas_utils import nearest_unequal_elements
+
+
+class TestNearestUnequalElements(ZiplineTestCase):
+
+    @parameter_space(tz=['UTC', 'US/Eastern'], __fail_fast=True)
+    def test_nearest_unequal_elements(self, tz):
+
+        dts = pd.to_datetime(
+            ['2014-01-01', '2014-01-05', '2014-01-06', '2014-01-09'],
+        ).tz_localize(tz)
+
+        t = lambda s: None if s is None else pd.Timestamp(s, tz=tz)
+
+        for dt, before, after in (('2013-12-30', None, '2014-01-01'),
+                                  ('2013-12-31', None, '2014-01-01'),
+                                  ('2014-01-01', None, '2014-01-05'),
+                                  ('2014-01-02', '2014-01-01', '2014-01-05'),
+                                  ('2014-01-03', '2014-01-01', '2014-01-05'),
+                                  ('2014-01-04', '2014-01-01', '2014-01-05'),
+                                  ('2014-01-05', '2014-01-01', '2014-01-06'),
+                                  ('2014-01-06', '2014-01-05', '2014-01-09'),
+                                  ('2014-01-07', '2014-01-06', '2014-01-09'),
+                                  ('2014-01-08', '2014-01-06', '2014-01-09'),
+                                  ('2014-01-09', '2014-01-06', None),
+                                  ('2014-01-10', '2014-01-09', None),
+                                  ('2014-01-11', '2014-01-09', None)):
+            computed = nearest_unequal_elements(dts, t(dt))
+            expected = (t(before), t(after))
+            self.assertEqual(computed, expected)
+
+    @parameter_space(tz=['UTC', 'US/Eastern'], __fail_fast=True)
+    def test_nearest_unequal_elements_short_dts(self, tz):
+
+        # Length 1.
+        dts = pd.to_datetime(['2014-01-01']).tz_localize(tz)
+        t = lambda s: None if s is None else pd.Timestamp(s, tz=tz)
+
+        for dt, before, after in (('2013-12-31', None, '2014-01-01'),
+                                  ('2014-01-01', None, None),
+                                  ('2014-01-02', '2014-01-01', None)):
+            computed = nearest_unequal_elements(dts, t(dt))
+            expected = (t(before), t(after))
+            self.assertEqual(computed, expected)
+
+        # Length 0
+        dts = pd.to_datetime([]).tz_localize(tz)
+        for dt, before, after in (('2013-12-31', None, None),
+                                  ('2014-01-01', None, None),
+                                  ('2014-01-02', None, None)):
+            computed = nearest_unequal_elements(dts, t(dt))
+            expected = (t(before), t(after))
+            self.assertEqual(computed, expected)
+
+    def test_nearest_unequal_bad_input(self):
+        with self.assertRaises(ValueError) as e:
+            nearest_unequal_elements(
+                pd.to_datetime(['2014', '2014']),
+                pd.Timestamp('2014'),
+            )
+
+        self.assertEqual(str(e.exception), 'dts must be unique')
+
+        with self.assertRaises(ValueError) as e:
+            nearest_unequal_elements(
+                pd.to_datetime(['2014', '2013']),
+                pd.Timestamp('2014'),
+            )
+
+        self.assertEqual(
+            str(e.exception),
+            'dts must be sorted in increasing order',
+        )

zipline/utils/pandas_utils.py

@@ -96,3 +96,53 @@ def mask_between_time(dts, start, end, include_start=True, include_end=True):
         left_op(start_micros, time_micros),
         right_op(time_micros, end_micros),
     )
+
+
+def nearest_unequal_elements(dts, dt):
+    """
+    Find values in ``dts`` closest but not equal to ``dt``.
+
+    Returns a pair of (last_before, first_after).
+
+    When ``dt`` is less than any element in ``dts``, ``last_before`` is None.
+    When ``dt`` is greater any element in ``dts``, ``first_after`` is None.
+
+    ``dts`` must be unique and sorted in increasing order.
+
+    Parameters
+    ----------
+    dts : pd.DatetimeIndex
+        Dates in which to search.
+    dt : pd.Timestamp
+        Date for which to find bounds.
+    """
+    if not dts.is_unique:
+        raise ValueError("dts must be unique")
+
+    if not dts.is_monotonic_increasing:
+        raise ValueError("dts must be sorted in increasing order")
+
+    if not len(dts):
+        return None, None
+
+    sortpos = dts.searchsorted(dt, side='left')
+    try:
+        sortval = dts[sortpos]
+    except IndexError:
+        # dt is greater than any value in the array.
+        return dts[-1], None
+
+    if dt < sortval:
+        lower_ix = sortpos - 1
+        upper_ix = sortpos
+    elif dt == sortval:
+        lower_ix = sortpos - 1
+        upper_ix = sortpos + 1
+    else:
+        lower_ix = sortpos
+        upper_ix = sortpos + 1
+
+    lower_value = dts[lower_ix] if lower_ix >= 0 else None
+    upper_value = dts[upper_ix] if upper_ix < len(dts) else None
+
+    return lower_value, upper_value