From d794bcb5b61c47e2a0cc6698c2e5a00baa1aa7e6 Mon Sep 17 00:00:00 2001 From: Bryant Moscon Date: Fri, 4 Nov 2016 08:50:45 -0400 Subject: [PATCH] Fix issue #275 --- CHANGES.md | 2 ++ arctic/chunkstore/date_chunker.py | 12 +++++------ tests/unit/chunkstore/test_date_chunker.py | 25 ++++++++++++++++++++++ 3 files changed, 33 insertions(+), 6 deletions(-) diff --git a/CHANGES.md b/CHANGES.md index 9f5bb89215f67..d6d967d4acc56 100644 --- 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 index 69c3de27185e4..8549bbad7debe 100644 --- 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 index 92eb98b69d51b..aba880ddd0efd 100644 --- 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)