Created StrategyLeg class from previous OptionContract. Option enums moved to separate module

backtester/option.py

@@ -1,31 +1,12 @@
+# Option Enum types
 from enum import Enum
 
+Type = Enum("Type", {"CALL": "call", "PUT": "put"})
+Direction = Enum("Direction", {"BUY": "ask", "SELL": "bid"})
 
-class OptionContract:
-    """Option contract data class"""
-
-    Type = Enum("Type", {"CALL": "call", "PUT": "put"})
-    Direction = Enum("Direction", "BUY SELL")
-
-    # Orders:
-    # BTO: Buy to Open
-    # BTC: Buy to Close
-    # STO: Sell to Open
-    # STC: Sell to Close
-    Order = Enum("Order", "BTO BTC STO STC")
-
-    def __init__(self,
-                 option_type=Type.CALL,
-                 direction=Direction.BUY,
-                 order=Order.BTO):
-        assert isinstance(option_type, OptionContract.Type)
-        assert isinstance(direction, OptionContract.Direction)
-        assert isinstance(order, OptionContract.Order)
-
-        self._store = {}
-        self._store["type"] = option_type
-        self._store["direction"] = direction
-        self._store["order"] = order
-
-    def __repr__(self):
-        return "Option({})".format(str(self._store))
+# Orders:
+# BTO: Buy to Open
+# BTC: Buy to Close
+# STO: Sell to Open
+# STC: Sell to Close
+Order = Enum("Order", "BTO BTC STO STC")

backtester/strategy/__init__.py

@@ -1 +1,2 @@
 from .strategy import Strategy
+from .strategy_leg import StrategyLeg

backtester/strategy/strategy_leg.py

@@ -0,0 +1,41 @@
+from backtester.option import Type, Direction
+from backtester.datahandler import Schema
+
+
+class StrategyLeg:
+    """Strategy Leg data class"""
+
+    def __init__(self, schema, option_type=Type.CALL, direction=Direction.BUY):
+        assert isinstance(schema, Schema)
+        assert isinstance(option_type, Type)
+        assert isinstance(direction, Direction)
+
+        self.schema = schema
+        self.type = option_type
+        self.direction = direction
+        self._entry_filter = self.schema.type == self.type.value
+        self._exit_filter = self.schema.type == self.type.value
+
+    @property
+    def entry_filter(self):
+        """Returns the entry filter"""
+        return self._entry_filter
+
+    @entry_filter.setter
+    def entry_filter(self, flt):
+        """Sets the entry filter"""
+        self._entry_filter = (self.schema.type == self.type.value) & flt
+
+    @property
+    def exit_filter(self):
+        """Returns the exit filter"""
+        return self._exit_filter
+
+    @exit_filter.setter
+    def exit_filter(self, flt):
+        """Sets the exit filter"""
+        self._exit_filter = (self.schema.type == self.type.value) & flt
+
+    def __repr__(self):
+        return "StrategyLeg(type={}, direction={}, entry_filter={}, exit_filter={})".format(
+            self.type, self.direction, self._entry_filter, self._exit_filter)