docstrings updates + pandas 0.19.0 compat changes

Author: bmoscon

.travis.yml

@@ -26,4 +26,4 @@ install:
   - pip install setuptools-git --upgrade
 script:
   - pip freeze
-  - python setup.py test
+  - python setup.py test --pytest-args=-v

arctic/chunkstore/chunkstore.py

@@ -308,10 +308,14 @@ def write(self, symbol, item, chunker=DateChunker(), **kwargs):
                                               {'$set': doc},
                                               upsert=True)
 
-    def __replace(self, old, new):
-        return new
-
     def __update(self, sym, item, combine_method=None, chunk_range=None):
+        '''
+        helper method used by update and append since they very closely
+        resemble eachother. Really differ only by the combine method.
+        append will combine existing date with new data (within a chunk),
+        whereas update will replace existing data with new data (within a
+        chunk).
+        '''
         if not isinstance(item, (DataFrame, Series)):
             raise Exception("Can only chunk DataFrames and Series")
 
@@ -426,7 +430,7 @@ def update(self, symbol, item, chunk_range=None, upsert=False, **kwargs):
                 raise Exception('Range must be inclusive of data')
             self.__update(sym, item, combine_method=self.serializer.combine, chunk_range=chunk_range)
         else:
-            self.__update(sym, item, combine_method=self.__replace, chunk_range=chunk_range)
+            self.__update(sym, item, combine_method=lambda old, new: new, chunk_range=chunk_range)
 
     def get_info(self, symbol):
         """
@@ -436,6 +440,10 @@ def get_info(self, symbol):
         ----------
         symbol: str
             the symbol for the given item in the DB
+
+        Returns
+        -------
+        dictionary
         """
         sym = self._get_symbol_info(symbol)
         if not sym:
@@ -461,6 +469,10 @@ def get_chunk_ranges(self, symbol, chunk_range=None, reverse=False):
             allows you to subset the chunks by range
         reverse: boolean
             return the chunk ranges in reverse order
+
+        Returns
+        -------
+        generator
         """
         sym = self._get_symbol_info(symbol)
         if not sym:
@@ -484,6 +496,10 @@ def iterator(self, symbol, chunk_range=None):
             the symbol for the given item in the DB
         chunk_range: None, or a range object
             allows you to subset the chunks by range
+
+        Returns
+        -------
+        generator
         """
         sym = self._get_symbol_info(symbol)
         if not sym:
@@ -505,6 +521,10 @@ def reverse_iterator(self, symbol, chunk_range=None):
             the symbol for the given item in the DB
         chunk_range: None, or a range object
             allows you to subset the chunks by range
+
+        Returns
+        -------
+        generator
         """
         sym = self._get_symbol_info(symbol)
         if not sym:

arctic/chunkstore/date_chunker.py

@@ -18,6 +18,9 @@ def to_chunks(self, df, chunk_size='D', **kwargs):
         """
         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')
@@ -26,8 +29,13 @@ def to_chunks(self, df, chunk_size='D', **kwargs):
         else:
             raise Exception("Data must be datetime indexed or have a column named 'date'")
 
-        for period, g in df.groupby(dates.to_period(chunk_size)):
-            start, end = period.start_time.to_pydatetime(warn=False), period.end_time.to_pydatetime(warn=False)
+        period_obj = dates.to_period(chunk_size)
+        period_obj_reduced = period_obj.drop_duplicates()
+        count = 0
+        for period, g in df.groupby(period_obj._data):
+            start = period_obj_reduced[count].start_time.to_pydatetime(warn=False)
+            end = period_obj_reduced[count].end_time.to_pydatetime(warn=False)
+            count += 1
             yield start, end, chunk_size, g
 
     def to_range(self, start, end):

tests/integration/chunkstore/test_chunkstore.py

@@ -311,7 +311,7 @@ def test_append_yearly(chunkstore_lib):
                                name='date'),
                    columns=['data'])
 
-    chunkstore_lib.write('chunkstore_test', df, chunk_size='Y')
+    chunkstore_lib.write('chunkstore_test', df, chunk_size='A')
     df2 = DataFrame(data=[4, 5, 6],
                     index=Index(data=[dt(2013, 1, 1),
                                       dt(2014, 1, 1),
@@ -563,7 +563,7 @@ def write_random_data(chunkstore_lib, name, month, days, securities, chunk_size=
     month: integer
     days: list of integers
     securities: list of integers
-    chunk_size: one of 'D', 'M', 'Y'
+    chunk_size: one of 'D', 'M', 'A'
     update: force update for each daily write
     append: force append for each daily write
     '''
@@ -597,7 +597,7 @@ def helper(chunkstore_lib, name, chunk_size):
         read_info = chunkstore_lib.read(name)
         assert_frame_equal(pd.concat([r, df]), read_info)
 
-    for chunk_size in ['D', 'M', 'Y']:
+    for chunk_size in ['D', 'M', 'A']:
         helper(chunkstore_lib, 'test_data_' + chunk_size, chunk_size)
 
 
@@ -614,7 +614,7 @@ def helper(chunkstore_lib, chunk_size, name, df, append):
 
         chunkstore_lib.update(name, append)
 
-        if chunk_size is not "Y":
+        if chunk_size is not 'A':
             assert_frame_equal(chunkstore_lib.read(name), pd.concat([df, append]))
         else:
             # chunksize is the entire DF, so we'll overwrite the whole thing
@@ -633,7 +633,7 @@ def helper(chunkstore_lib, chunk_size, name, df, append):
 
     append = pd.concat(append)
 
-    for chunk_size in ['D', 'M', 'Y']:
+    for chunk_size in ['D', 'M', 'A']:
         helper(chunkstore_lib, chunk_size, 'test_monthly_' + chunk_size, df, append)
 
 
@@ -870,7 +870,7 @@ def test_yearly_series(chunkstore_lib):
                             name='date'),
                 name='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_series_equal(df, ret)
 
@@ -1067,8 +1067,8 @@ def test_size_chunk_append(chunkstore_lib):
 
 
 def test_delete_range_segment(chunkstore_lib):
-    df = DataFrame(data={'data': np.random.randint(0, 100, size=20000000),
-                         'date': [dt(2016, 1, 1)] * 20000000})
+    df = DataFrame(data={'data': np.random.randint(0, 100, size=7000000),
+                         'date': [dt(2016, 1, 1)] * 7000000})
     dg = DataFrame(data={'data': np.random.randint(0, 100, size=100),
                          'date': [dt(2016, 1, 2)] * 100})
     chunkstore_lib.write('test_df', pd.concat([df, dg], ignore_index=True), chunk_size='M')