Skip to content

Commit d8cfa15

Browse files
author
Ernesto Perez
committed
ENH: Expose SQLiteAdjustmentReader.get_df_from_table method.
1 parent 660b4e8 commit d8cfa15

File tree

1 file changed

+29
-23
lines changed

1 file changed

+29
-23
lines changed

zipline/data/adjustments.py

Lines changed: 29 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -197,33 +197,39 @@ def unpack_db_to_component_dfs(self, convert_dates=False):
197197
version of the table, where all date columns have been coerced back
198198
from int to datetime.
199199
"""
200+
return {
201+
t_name: self.get_df_from_table(t_name, convert_dates)
202+
for t_name in self._datetime_int_cols
203+
}
200204

201-
def _get_df_from_table(table_name, date_cols):
202-
203-
# Dates are stored in second resolution as ints in adj.db tables.
204-
# Need to specifically convert them as UTC, not local time.
205-
kwargs = (
206-
{'parse_dates': {col: {'unit': 's', 'utc': True}
207-
for col in date_cols}
208-
}
209-
if convert_dates
210-
else {}
205+
def get_df_from_table(self, table_name, convert_dates=False):
206+
try:
207+
date_cols = self._datetime_int_cols[table_name]
208+
except KeyError:
209+
raise ValueError(
210+
"Requested table %s not found.\n"
211+
"Available tables: %s\n" % (
212+
table_name,
213+
self._datetime_int_cols.keys(),
214+
)
211215
)
212216

213-
return pd.read_sql(
214-
'select * from "{}"'.format(table_name),
215-
self.conn,
216-
index_col='index',
217-
**kwargs
218-
).rename_axis(None)
217+
# Dates are stored in second resolution as ints in adj.db tables.
218+
# Need to specifically convert them as UTC, not local time.
219+
kwargs = (
220+
{'parse_dates': {col: {'unit': 's', 'utc': True}
221+
for col in date_cols}
222+
}
223+
if convert_dates
224+
else {}
225+
)
219226

220-
return {
221-
t_name: _get_df_from_table(
222-
t_name,
223-
date_cols
224-
)
225-
for t_name, date_cols in self._datetime_int_cols.items()
226-
}
227+
return pd.read_sql(
228+
'select * from "{}"'.format(table_name),
229+
self.conn,
230+
index_col='index',
231+
**kwargs
232+
).rename_axis(None)
227233

228234

229235
class SQLiteAdjustmentWriter(object):

0 commit comments

Comments
 (0)