ObjectStoreHelpersRegressionAlgorithm.py

# QUANTCONNECT.COM - Democratizing Finance, Empowering Individuals.
# Lean Algorithmic Trading Engine v2.0. Copyright 2014 QuantConnect Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# 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 AlgorithmImports import *
from System.Collections.Generic import KeyNotFoundException, List
from decimal import Decimal

### <summary>
### Regression algorithm asserting the ObjectStore helpers: save_text and read_text are aliases of save and read,
### reading a missing key lists the available keys in the error, and the tolerant save_json/read_json and
### save_dataframe/read_dataframe helpers handle datetime/date/Decimal/Symbol and data frames out of the box.
### </summary>
class ObjectStoreHelpersRegressionAlgorithm(QCAlgorithm):

    def initialize(self):
        self.set_start_date(2013, 10, 7)
        self.set_end_date(2013, 10, 11)

        self.set_benchmark(lambda x: 0)

        self.spy = self.add_equity("SPY", Resolution.DAILY).symbol

    def on_end_of_algorithm(self):
        # save_json tolerates datetime/date/Decimal/numpy/NaN values, .NET objects and non-string dictionary keys
        trade_log = {
            "time": datetime(2013, 10, 11, 16, 0, 0),
            "date": date(2013, 10, 11),
            "qty": Decimal("1.5"),
            "symbol": self.spy,
            "fills": [{"price": Decimal("167.42"), self.spy: 10}],
            "fill_quantities": List[int]([10, -10]),
            "returns": np.array([0.01, np.nan]),
            "bar_times": pd.DatetimeIndex([datetime(2013, 10, 10), datetime(2013, 10, 11)]).values
        }
        if not self.object_store.save_json("helpers/trade_log.json", trade_log):
            raise AssertionError("Expected save_json to succeed")
        data = self.object_store.read_json("helpers/trade_log.json")
        expected = {
            "time": "2013-10-11T16:00:00",
            "date": "2013-10-11",
            "qty": 1.5,
            "symbol": {"value": "SPY", "id": str(self.spy.id), "permtick": "SPY"},
            "fills": [{"price": 167.42, "SPY": 10}],
            "fill_quantities": [10, -10],
            "returns": [0.01, None],
            "bar_times": ["2013-10-10T00:00:00", "2013-10-11T00:00:00"]
        }
        if data != expected:
            raise AssertionError(f"Unexpected read_json round trip result: {data}")
        # the symbol id and the ticker key both resolve back to the added security through the symbol cache
        if self.symbol(data["symbol"]["id"]) != self.spy or self.symbol("SPY") != self.spy:
            raise AssertionError(f"Expected the stored symbol to resolve to {self.spy}, got: {data['symbol']}")

        # save_json refuses what it cannot serialize faithfully instead of storing its string
        error = None
        try:
            self.object_store.save_json("helpers/rejected.json", {"history": pd.Series([1.0, 2.0])})
        except TypeError as e:
            error = str(e)
        if error is None or "save_dataframe" not in error:
            raise AssertionError(f"Expected save_json to reject a pandas Series pointing to save_dataframe, got: '{error}'")

        # read_json throws on a missing key like every other read method
        error = None
        try:
            self.object_store.read_json("helpers/missing.json")
        except KeyNotFoundException as e:
            error = str(e)
        if error is None:
            raise AssertionError("Expected read_json of a missing key to throw")

        # save_dataframe stores a pandas DataFrame as CSV and read_dataframe restores it, time index included
        frame = pd.DataFrame({"close": [167.42, 168.0]}, index=pd.to_datetime(["2013-10-10", "2013-10-11"]))
        if not self.object_store.save_dataframe("helpers/history.csv", frame):
            raise AssertionError("Expected save_dataframe to succeed")
        csv = self.object_store.read("helpers/history.csv")
        if "close" not in csv or "2013-10-10" not in csv:
            raise AssertionError(f"Unexpected save_dataframe content: {csv}")
        restored = self.object_store.read_dataframe("helpers/history.csv")
        if not restored.equals(frame):
            raise AssertionError(f"Unexpected read_dataframe round trip result: {restored}")

        # save_text/read_text are aliases of save/read
        if not self.object_store.save_text("helpers/report.txt", "The strategy went up"):
            raise AssertionError("Expected save_text to succeed")
        if self.object_store.read_text("helpers/report.txt") != "The strategy went up":
            raise AssertionError("Expected the save_text/read_text round trip to succeed")

        # reading a missing key lists the available keys in the error
        error = None
        try:
            self.object_store.read_bytes("helpers/missing.json")
        except BaseException as e:
            error = str(e)
        if error is None or "Keys: [" not in error or "'helpers/report.txt'" not in error:
            raise AssertionError(f"Expected the available keys in the missing key error, got: '{error}'")

        for key in ["helpers/trade_log.json", "helpers/history.csv", "helpers/report.txt"]:
            self.object_store.delete(key)