From 5025101d37bd2b44bae1d6f71e0f332255d0d10b Mon Sep 17 00:00:00 2001 From: Kathryn Glowinski Date: Tue, 27 Dec 2016 13:44:17 -0500 Subject: [PATCH] Adjustments to Component Dfs (#1620) * ENH: SQLiteAdjustmentReader can return DF versions of tables. --- .../pipeline/test_us_equity_pricing_loader.py | 53 +++++++++++++++ zipline/data/us_equity_pricing.py | 68 +++++++++++++++++-- 2 files changed, 115 insertions(+), 6 deletions(-) diff --git a/tests/pipeline/test_us_equity_pricing_loader.py b/tests/pipeline/test_us_equity_pricing_loader.py index 527bec61..39a887c2 100644 --- a/tests/pipeline/test_us_equity_pricing_loader.py +++ b/tests/pipeline/test_us_equity_pricing_loader.py @@ -15,6 +15,7 @@ """ Tests for USEquityPricingLoader and related classes. """ +from nose_parameterized import parameterized from numpy import ( arange, datetime64, @@ -32,6 +33,7 @@ from pandas import ( Int64Index, Timestamp, ) +from pandas.util.testing import assert_frame_equal from toolz.curried.operator import getitem from zipline.lib.adjustment import Float64Multiply @@ -395,6 +397,57 @@ class USEquityPricingLoaderTestCase(WithAdjustmentReader, self.assertEqual(adj.last_col, expected.last_col) assert_allclose(adj.value, expected.value) + @parameterized([(True,), (False,)]) + def test_load_adjustments_to_df(self, convert_dts): + reader = self.adjustment_reader + adjustment_dfs = reader.unpack_db_to_component_dfs( + convert_dates=convert_dts + ) + + name_and_raw = ( + ('splits', SPLITS), + ('mergers', MERGERS), + ('dividends', DIVIDENDS_EXPECTED) + ) + + def create_expected_table(df, name): + expected_df = df.copy() + + if convert_dts: + for colname in reader._datetime_int_cols[name]: + expected_df[colname] = expected_df[colname].astype( + 'datetime64[s]' + ) + + return expected_df + + def create_expected_div_table(df, name): + expected_df = df.copy() + + if not convert_dts: + for colname in reader._datetime_int_cols[name]: + expected_df[colname] = expected_df[colname].astype( + 'datetime64[s]' + ).astype(int) + + return expected_df + + for action_name, raw_tbl in name_and_raw: + + exp = create_expected_table(raw_tbl, action_name) + assert_frame_equal( + adjustment_dfs[action_name], + exp + ) + + # DIVIDENDS is in the opposite form from the rest of the dataframes, so + # needs to be converted separately. + div_name = 'dividend_payouts' + assert_frame_equal( + adjustment_dfs[div_name], + create_expected_div_table(DIVIDENDS, div_name) + ) + def test_read_no_adjustments(self): adjustment_reader = NullAdjustmentReader() columns = [USEquityPricing.close, USEquityPricing.volume] diff --git a/zipline/data/us_equity_pricing.py b/zipline/data/us_equity_pricing.py index 5165e19c..94a57e1f 100644 --- a/zipline/data/us_equity_pricing.py +++ b/zipline/data/us_equity_pricing.py @@ -36,18 +36,19 @@ from numpy import ( uint32, ) from pandas import ( - isnull, DataFrame, - read_csv, - Timestamp, + DatetimeIndex, + isnull, NaT, - DatetimeIndex + read_csv, + read_sql, + Timestamp, ) from pandas.tslib import iNaT from six import ( iteritems, - viewkeys, string_types, + viewkeys, ) from zipline.data.session_bars import SessionBarReader @@ -60,8 +61,8 @@ from zipline.utils.calendars import get_calendar from zipline.utils.functional import apply from zipline.utils.preprocess import call from zipline.utils.input_validation import ( - preprocess, expect_element, + preprocess, verify_indices_all_unique, ) from zipline.utils.sqlite_utils import group_into_chunks, coerce_string_to_conn @@ -1250,6 +1251,18 @@ class SQLiteAdjustmentReader(object): def __init__(self, conn): self.conn = conn + # Given the tables in the adjustments.db file, dict which knows which + # col names contain dates that have been coerced into ints. + self._datetime_int_cols = { + 'dividend_payouts': ('declared_date', 'ex_date', 'pay_date', + 'record_date'), + 'dividends': ('effective_date',), + 'mergers': ('effective_date',), + 'splits': ('effective_date',), + 'stock_dividend_payouts': ('declared_date', 'ex_date', 'pay_date', + 'record_date') + } + def load_adjustments(self, columns, dates, assets): return load_adjustments_from_sqlite( self.conn, @@ -1316,3 +1329,46 @@ class SQLiteAdjustmentReader(object): c.close() return stock_divs + + def unpack_db_to_component_dfs(self, convert_dates=False): + """Returns the set of known tables in the adjustments file in DataFrame + form. + + Parameters + ---------- + convert_dates : bool, optional + By default, dates are returned in seconds since EPOCH. If + convert_dates is True, all ints in date columns will be converted + to datetimes. + + Returns + ------- + dfs : dict{str->DataFrame} + Dictionary which maps table name to the corresponding DataFrame + version of the table, where all date columns have been coerced back + from int to datetime. + """ + + def _get_df_from_table(table_name, date_cols): + + kwargs = ( + {'parse_dates': {col: 's' for col in date_cols}} + if convert_dates + else {} + ) + + # Dates are stored in second resolution as ints in adj.db tables. + return read_sql( + 'select * from "{}"'.format(table_name), + self.conn, + index_col='index', + **kwargs + ).rename_axis(None) + + return { + t_name: _get_df_from_table( + t_name, + date_cols + ) + for t_name, date_cols in self._datetime_int_cols.items() + }