diff --git a/docs/source/whatsnew/0.8.4.txt b/docs/source/whatsnew/0.8.4.txt index 2fa84f31..7dfd070a 100644 --- a/docs/source/whatsnew/0.8.4.txt +++ b/docs/source/whatsnew/0.8.4.txt @@ -63,6 +63,8 @@ Enhancements subclasses inherit all of the columns from the parent. These columns will be new sentinels so you can register them a custom loader (:issue:`924`). +* Added :func:`~zipline.utils.input_validation.coerce`. + Experimental Features ~~~~~~~~~~~~~~~~~~~~~ diff --git a/zipline/data/us_equity_pricing.py b/zipline/data/us_equity_pricing.py index 236dae5b..dcd9b135 100644 --- a/zipline/data/us_equity_pricing.py +++ b/zipline/data/us_equity_pricing.py @@ -23,6 +23,7 @@ import sqlite3 from bcolz import ( carray, ctable, + open as open_ctable, ) from click import progressbar from numpy import ( @@ -45,10 +46,11 @@ from pandas import ( ) from six import ( iteritems, - string_types, with_metaclass, ) +from zipline.utils.input_validation import coerce_string, preprocess + from ._equities import _compute_row_slices, _read_bcolz_data from ._adjustments import load_adjustments_from_sqlite @@ -363,9 +365,8 @@ class BcolzDailyBarReader(object): We use calendar_offset and calendar to orient loaded blocks within a range of queried dates. """ + @preprocess(table=coerce_string(open_ctable, mode='r')) def __init__(self, table): - if isinstance(table, string_types): - table = ctable(rootdir=table, mode='r') self._table = table self._calendar = DatetimeIndex(table.attrs['calendar'], tz='UTC') @@ -910,9 +911,8 @@ class SQLiteAdjustmentReader(object): Connection from which to load data. """ + @preprocess(conn=coerce_string(sqlite3.connect)) def __init__(self, conn): - if isinstance(conn, str): - conn = sqlite3.connect(conn) self.conn = conn def load_adjustments(self, columns, dates, assets): diff --git a/zipline/utils/input_validation.py b/zipline/utils/input_validation.py index 948f8108..4d7256a5 100644 --- a/zipline/utils/input_validation.py +++ b/zipline/utils/input_validation.py @@ -11,6 +11,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +from functools import partial from operator import attrgetter from numpy import dtype @@ -283,6 +284,46 @@ def expect_element(*_pos, **named): return preprocess(**valmap(_expect_element, named)) +def coerce(from_, to, **to_kwargs): + """ + A preprocessing decorator that coerces inputs of a given type by passing + them to a callable. + + Parameters + ---------- + from : type or tuple or types + Inputs types on which to call ``to``. + to : function + Coercion function to call on inputs. + **to_kwargs + Additional keywords to forward to every call to ``to``. + + Usage + ----- + >>> @preprocess(x=coerce(float, int), y=coerce(float, int)) + ... def floordiff(x, y): + ... return x - y + ... + >>> floordiff(3.2, 2.5) + 1 + + >>> @preprocess(x=coerce(str, int, base=2), y=coerce(str, int, base=2)) + ... def add_binary_strings(x, y): + ... return bin(x + y)[2:] + ... + >>> add_binary_strings('101', '001') + '110' + """ + def preprocessor(func, argname, arg): + if isinstance(arg, from_): + return to(arg, **to_kwargs) + return arg + return preprocessor + + +coerce_string = partial(coerce, string_types) + + def _expect_element(collection): template = ( "%(funcname)s() expected a value in {collection} "