Fix issue pandas-dev#283

Author: bmoscon

CHANGES.md

@@ -1,5 +1,9 @@
 ## Changelog
 
+### 1.34
+
+  * Feature: #283 Support for all pandas frequency strings in ChunkStore DateChunker
+
 ### 1.33 (2016-11-07)
 
   * Feature:    #275 Tuple range object support in DateChunker

arctic/chunkstore/date_chunker.py

@@ -11,17 +11,17 @@ def to_chunks(self, df, chunk_size='D', **kwargs):
         """
         chunks the dataframe/series by dates
 
-        returns
+        Parameters
+        ----------
+        df: pandas dataframe or series
+        chunk_size: str
+            any valid Pandas frequency string
+
+        Returns
         -------
         generator that produces tuples: (start date, end date,
                   chunk_size, dataframe/series)
         """
-        if chunk_size not in ('D', 'M', 'Y', 'A'):
-            raise Exception("Chunk size must be one of D, M, Y, A")
-
-        if chunk_size == 'Y':
-            chunk_size = 'A'
-
         if 'date' in df.index.names:
             dates = df.index.get_level_values('date')
         elif 'date' in df.columns:

tests/integration/chunkstore/test_chunkstore.py

@@ -258,7 +258,7 @@ def test_yearly_df(chunkstore_lib):
                                name='date'),
                    columns=['data'])
 
-    chunkstore_lib.write('chunkstore_test', df, chunk_size='Y')
+    chunkstore_lib.write('chunkstore_test', df, chunk_size='A')
     ret = chunkstore_lib.read('chunkstore_test', chunk_range=DateRange(dt(2016, 1, 1), dt(2016, 3, 3)))
     assert_frame_equal(df, ret)
 
@@ -1172,3 +1172,18 @@ def test_unnamed_colums(chunkstore_lib):
     with pytest.raises(Exception) as e:
         chunkstore_lib.write('test_df', df, chunk_size='D')
     assert('must be named' in str(e))
+
+
+def test_quarterly_data(chunkstore_lib):
+    df = DataFrame(data={'data': np.random.randint(0, 100, size=366)},
+                   index=pd.date_range('2016-01-01', '2016-12-31'))
+    df.index.name = 'date'
+
+    chunkstore_lib.write('quarterly', df, chunk_size='Q')
+    assert_frame_equal(df, chunkstore_lib.read('quarterly'))
+    assert(len(chunkstore_lib.read('quarterly', chunk_range=(None, '2016-01-05'))) == 5)
+    count = 0
+    for _ in chunkstore_lib._collection.find({SYMBOL: 'quarterly'}, sort=[(START, pymongo.ASCENDING)],):
+        count += 1
+
+    assert(count == 4)

tests/unit/chunkstore/test_date_chunker.py

@@ -74,14 +74,16 @@ def test_date_filter_with_pd_date_range():
 def test_to_chunks_exceptions():
     df = DataFrame(data={'data': [1, 2, 3]})
     c = DateChunker()
-    with pytest.raises(Exception) as e:
-        six.next(c.to_chunks(None, None))
-    assert('Chunk size' in str(e))
 
     with pytest.raises(Exception) as e:
         six.next(c.to_chunks(df, 'D'))
     assert('datetime indexed' in str(e))
 
+    df.columns = ['date']
+    with pytest.raises(Exception) as e:
+        six.next(c.to_chunks(df, 'ZSDFG'))
+    assert('Unknown freqstr' in str(e) or 'Invalid frequency' in str(e))
+
 
 def test_exclude():
     c = DateChunker()