Add tagging, object typing, fix pep8

Author: qacollective

backtesting/backtesting.py

@@ -47,6 +47,7 @@ class Strategy(metaclass=ABCMeta):
     `backtesting.backtesting.Strategy.next` to define
     your own strategy.
     """
+
     def __init__(self, broker, data, params):
         self._indicators = []
         self._broker: _Broker = broker
@@ -193,30 +194,32 @@ def buy(self, *,
             limit: float = None,
             stop: float = None,
             sl: float = None,
-            tp: float = None):
+            tp: float = None,
+            tag: object = None):
         """
         Place a new long order. For explanation of parameters, see `Order` and its properties.
 
         See also `Strategy.sell()`.
         """
         assert 0 < size < 1 or round(size) == size, \
             "size must be a positive fraction of equity, or a positive whole number of units"
-        return self._broker.new_order(size, limit, stop, sl, tp)
+        return self._broker.new_order(size, limit, stop, sl, tp, tag)
 
     def sell(self, *,
              size: float = 1 - sys.float_info.epsilon,
              limit: float = None,
              stop: float = None,
              sl: float = None,
-             tp: float = None):
+             tp: float = None,
+             tag: object = None):
         """
         Place a new short order. For explanation of parameters, see `Order` and its properties.
 
         See also `Strategy.buy()`.
         """
         assert 0 < size < 1 or round(size) == size, \
             "size must be a positive fraction of equity, or a positive whole number of units"
-        return self._broker.new_order(-size, limit, stop, sl, tp)
+        return self._broker.new_order(-size, limit, stop, sl, tp, tag)
 
     @property
     def equity(self) -> float:
@@ -277,6 +280,7 @@ class _Orders(tuple):
     """
     TODO: remove this class. Only for deprecation.
     """
+
     def cancel(self):
         """Cancel all non-contingent (i.e. SL/TP) orders."""
         for order in self:
@@ -304,6 +308,7 @@ class Position:
         if self.position:
             ...  # we have a position, either long or short
     """
+
     def __init__(self, broker: '_Broker'):
         self.__broker = broker
 
@@ -368,13 +373,15 @@ class Order:
     [filled]: https://www.investopedia.com/terms/f/fill.asp
     [Good 'Til Canceled]: https://www.investopedia.com/terms/g/gtc.asp
     """
+
     def __init__(self, broker: '_Broker',
                  size: float,
                  limit_price: float = None,
                  stop_price: float = None,
                  sl_price: float = None,
                  tp_price: float = None,
-                 parent_trade: 'Trade' = None):
+                 parent_trade: 'Trade' = None,
+                 tag: object = None):
         self.__broker = broker
         assert size != 0
         self.__size = size
@@ -383,6 +390,7 @@ def __init__(self, broker: '_Broker',
         self.__sl_price = sl_price
         self.__tp_price = tp_price
         self.__parent_trade = parent_trade
+        self.__tag = tag
 
     def _replace(self, **kwargs):
         for k, v in kwargs.items():
@@ -398,6 +406,7 @@ def __repr__(self):
                                                  ('sl', self.__sl_price),
                                                  ('tp', self.__tp_price),
                                                  ('contingent', self.is_contingent),
+                                                 ('tag', self.__tag),
                                              ) if value is not None))
 
     def cancel(self):
@@ -468,6 +477,15 @@ def tp(self) -> Optional[float]:
     def parent_trade(self):
         return self.__parent_trade
 
+    @property
+    def tag(self) -> Optional[object]:
+        """
+        An attribute which, if set, persists to enable tracking of this order
+        by an external identifier if it becomes a trade in `Strategy.trades`
+        and when closed in `Strategy.closed_trades`.
+        """
+        return self.__tag
+
     __pdoc__['Order.parent_trade'] = False
 
     # Extra properties
@@ -502,7 +520,8 @@ class Trade:
     When an `Order` is filled, it results in an active `Trade`.
     Find active trades in `Strategy.trades` and closed, settled trades in `Strategy.closed_trades`.
     """
-    def __init__(self, broker: '_Broker', size: int, entry_price: float, entry_bar):
+
+    def __init__(self, broker: '_Broker', size: int, entry_price: float, entry_bar, tag: object):
         self.__broker = broker
         self.__size = size
         self.__entry_price = entry_price
@@ -511,10 +530,12 @@ def __init__(self, broker: '_Broker', size: int, entry_price: float, entry_bar):
         self.__exit_bar: Optional[int] = None
         self.__sl_order: Optional[Order] = None
         self.__tp_order: Optional[Order] = None
+        self.__tag = tag
 
     def __repr__(self):
         return f'<Trade size={self.__size} time={self.__entry_bar}-{self.__exit_bar or ""} ' \
-               f'price={self.__entry_price}-{self.__exit_price or ""} pl={self.pl:.0f}>'
+               f'price={self.__entry_price}-{self.__exit_price or ""} pl={self.pl:.0f}' \
+               f'{" tag="+str(self.__tag) if self.__tag is not None else ""}>'
 
     def _replace(self, **kwargs):
         for k, v in kwargs.items():
@@ -528,7 +549,7 @@ def close(self, portion: float = 1.):
         """Place new `Order` to close `portion` of the trade at next market price."""
         assert 0 < portion <= 1, "portion must be a fraction between 0 and 1"
         size = copysign(max(1, round(abs(self.__size) * portion)), -self.__size)
-        order = Order(self.__broker, size, parent_trade=self)
+        order = Order(self.__broker, size, parent_trade=self, tag=self.__tag)
         self.__broker.orders.insert(0, order)
 
     # Fields getters
@@ -561,6 +582,15 @@ def exit_bar(self) -> Optional[int]:
         """
         return self.__exit_bar
 
+    @property
+    def tag(self) -> Optional[object]:
+        """
+        A tag attribute optionally set when placing an order with
+        `Strategy.buy()` or `Strategy.sell()`.
+        See `Order.tag`.
+        """
+        return self.__tag
+
     @property
     def _sl_order(self):
         return self.__sl_order
@@ -652,7 +682,7 @@ def __set_contingent(self, type, price):
             order.cancel()
         if price:
             kwargs = dict(stop=price) if type == 'sl' else dict(limit=price)
-            order = self.__broker.new_order(-self.size, trade=self, **kwargs)
+            order = self.__broker.new_order(-self.size, trade=self, tag=self.tag, **kwargs)
             setattr(self, attr, order)
 
 
@@ -685,6 +715,7 @@ def new_order(self,
                   stop: float = None,
                   sl: float = None,
                   tp: float = None,
+                  tag: object = None,
                   *,
                   trade: Trade = None):
         """
@@ -710,7 +741,7 @@ def new_order(self,
                     "Short orders require: "
                     f"TP ({tp}) < LIMIT ({limit or stop or adjusted_price}) < SL ({sl})")
 
-        order = Order(self, size, limit, stop, sl, tp, trade)
+        order = Order(self, size, limit, stop, sl, tp, trade, tag)
         # Put the new order in the order queue,
         # inserting SL/TP/trade-closing orders in-front
         if trade:
@@ -890,7 +921,12 @@ def _process_orders(self):
 
             # Open a new trade
             if need_size:
-                self._open_trade(adjusted_price, need_size, order.sl, order.tp, time_index)
+                self._open_trade(adjusted_price,
+                                 need_size,
+                                 order.sl,
+                                 order.tp,
+                                 time_index,
+                                 order.tag)
 
                 # We need to reprocess the SL/TP orders newly added to the queue.
                 # This allows e.g. SL hitting in the same bar the order was open.
@@ -948,8 +984,8 @@ def _close_trade(self, trade: Trade, price: float, time_index: int):
         self.closed_trades.append(trade._replace(exit_price=price, exit_bar=time_index))
         self._cash += trade.pl
 
-    def _open_trade(self, price: float, size: int, sl: float, tp: float, time_index: int):
-        trade = Trade(self, size, price, time_index)
+    def _open_trade(self, price: float, size: int, sl: float, tp: float, time_index: int, tag: object):
+        trade = Trade(self, size, price, time_index, tag)
         self.trades.append(trade)
         # Create SL/TP (bracket) orders.
         # Make sure SL order is created first so it gets adversarially processed before TP order
@@ -971,6 +1007,7 @@ class Backtest:
     instance, or `backtesting.backtesting.Backtest.optimize` to
     optimize it.
     """
+
     def __init__(self,
                  data: pd.DataFrame,
                  strategy: Type[Strategy],