From 1b4846b652b3bcc2274eb37a95d2639a984ae4eb Mon Sep 17 00:00:00 2001 From: Juan Pablo Amoroso Date: Fri, 31 May 2019 16:11:08 -0300 Subject: [PATCH] Added `apply_filter` method to datahandler. Fixed schema updates for `__setitem__` calls --- backtester/datahandler/historical_options_data.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/backtester/datahandler/historical_options_data.py b/backtester/datahandler/historical_options_data.py index 0bc9fc0..091dce5 100644 --- a/backtester/datahandler/historical_options_data.py +++ b/backtester/datahandler/historical_options_data.py @@ -15,16 +15,23 @@ class HistoricalOptionsData: columns = self._data.columns assert all((col in columns for _key, col in self.schema)) + self.index = self._data.index self._data["dte"] = (self._data["expiration"] - self._data["quotedate"]).dt.days self.schema.update({"dte": "dte"}) + def apply_filter(self, f): + """Apply Filter `f` to the data. Returns a `pd.DataFrame` with the filtered rows.""" + return self._data.query(f.query) + def __getitem__(self, item): - key = self.schema[item].mapping + key = self.schema[item] return self._data[key] - def __setitem__(self, item, value): - self._data[item] = value + def __setitem__(self, key, value): + self._data[key] = value + if key not in self.schema: + self.schema.update({key: key}) def __repr__(self): return self._data.__repr__()