Fix issue pandas-dev#275

dan-nadler · Nov 4, 2016 · d794bcb · d794bcb
1 parent 5794407
commit d794bcb
Show file tree

Hide file tree

Showing 3 changed files with 33 additions and 6 deletions.
diff --git a/CHANGES.md b/CHANGES.md
@@ -2,6 +2,8 @@
 
 ### 1.33
 
+  * Feature:    #275 Tuple range object support in DateChunker
+  * Bugfix:     #273 Duplicate columns breaking serializer
   * Feature:    #267 Tickstore.delete returns deleted data
   * Dependency: #266 Remove pytest-dbfixtures in favor of pytest-server-fixtures
 

diff --git a/arctic/chunkstore/date_chunker.py b/arctic/chunkstore/date_chunker.py
@@ -71,8 +71,8 @@ def to_mongo(self, range_obj):
         -------
         dict
         """
-        if isinstance(range_obj, pd.DatetimeIndex):
-            range_obj = DateRange(range_obj.min(), range_obj.max())
+        if isinstance(range_obj, (pd.DatetimeIndex, tuple)):
+            range_obj = DateRange(range_obj[0], range_obj[-1])
         if range_obj.start and range_obj.end:
             return {'$and': [{START: {'$lte': range_obj.end}}, {END: {'$gte': range_obj.start}}]}
         elif range_obj.start:
@@ -95,8 +95,8 @@ def filter(self, data, range_obj):
         -------
         data, filtered by range_obj
         """
-        if isinstance(range_obj, pd.DatetimeIndex):
-            range_obj = DateRange(range_obj.min(), range_obj.max())
+        if isinstance(range_obj, (pd.DatetimeIndex, tuple)):
+            range_obj = DateRange(range_obj[0], range_obj[-1])
         if 'date' in data.index.names:
             return data[range_obj.start:range_obj.end]
         elif 'date' in data.columns:
@@ -119,8 +119,8 @@ def exclude(self, data, range_obj):
         -------
         data, filtered by range_obj
         """
-        if isinstance(range_obj, pd.DatetimeIndex):
-            range_obj = DateRange(range_obj.min(), range_obj.max())
+        if isinstance(range_obj, (pd.DatetimeIndex, tuple)):
+            range_obj = DateRange(range_obj[0], range_obj[-1])
         if 'date' in data.index.names:
             return data[(data.index.get_level_values('date') < range_obj.start) | (data.index.get_level_values('date') > range_obj.end)]
         elif 'date' in data.columns:

diff --git a/tests/unit/chunkstore/test_date_chunker.py b/tests/unit/chunkstore/test_date_chunker.py
@@ -111,3 +111,28 @@ def test_exclude_no_index():
 
     assert(c.exclude(df, DateRange(dt(2016, 1, 1), dt(2016, 1, 1))).equals(c.exclude(df, pd.date_range(dt(2016, 1, 1), dt(2016, 1, 1)))))
     assert(c.exclude(df2, None).equals(df2))
+
+
+def test_with_tuples():
+    c = DateChunker()
+    df = DataFrame(data={'data': [1, 2, 3],
+                         'date': [dt(2016, 1, 1),
+                                  dt(2016, 1, 2),
+                                  dt(2016, 1, 3)]
+                         }
+                   )
+
+    # OPEN - CLOSED
+    assert_frame_equal(c.filter(df, (None, dt(2016, 1, 3))), df)
+    # CLOSED - OPEN
+    assert_frame_equal(c.filter(df, (dt(2016, 1, 1), None)), df)
+    # OPEN - OPEN
+    assert_frame_equal(c.filter(df, (None, None)), df)
+    # CLOSED - OPEN (far before data range)
+    assert_frame_equal(c.filter(df, (dt(2000, 1, 1), None)), df)
+    # CLOSED - OPEN (far after range)
+    assert(c.filter(df, (dt(2020, 1, 2), None)).empty)
+    # OPEN - CLOSED
+    assert_frame_equal(c.filter(df, (None, dt(2020, 1, 1))), df)
+    # CLOSED - CLOSED (after range)
+    assert(c.filter(df, (dt(2017, 1, 1), dt(2018, 1, 1))).empty)