bug: backtest.py bug fixes

python/backtest.py

@@ -1,8 +1,15 @@
-import backtestingfx as _rust
+from typing import Any
+
+import backtestingfx as _rust  # type: ignore
 import pandas as pd
 
 
 class Strategy:
+    def __init__(self):
+        self._bars: Any = None
+        self._bar: Any = None
+        self._broker: Any = None
+
     def init(self):
         pass
 
@@ -10,10 +17,14 @@ class Strategy:
         pass
 
     def buy(self, lot_size, stop_loss=None, take_profit=None):
-        self._broker.buy(self._bar.close, lot_size, self._bar.timestamp, stop_loss, take_profit)
+        self._broker.buy(
+            self._bar.close, lot_size, self._bar.timestamp, stop_loss, take_profit
+        )
 
     def sell(self, lot_size, stop_loss=None, take_profit=None):
-        self._broker.sell(self._bar.close, lot_size, self._bar.timestamp, stop_loss, take_profit)
+        self._broker.sell(
+            self._bar.close, lot_size, self._bar.timestamp, stop_loss, take_profit
+        )
 
     def close_all(self):
         self._broker.close_all(self._bar.close, self._bar.timestamp)
@@ -27,7 +38,7 @@ class _Adapter:
         self._strategy = strategy
 
     def init(self, bars):
-        self._strategy.bars = bars
+        self._strategy._bars = bars
         self._strategy.init()
 
     def next(self, bar, broker):
@@ -50,20 +61,22 @@ class Backtest:
             if isinstance(idx, pd.Timestamp):
                 ts = int(idx.timestamp())
             else:
-                ts = int(pd.Timestamp(row["timestamp"]).timestamp())
+                ts = int(pd.Timestamp(row["timestamp"]).timestamp())  # type: ignore
 
-            bars.append(_rust.Bar(
-                timestamp=ts,
-                open=float(row["open"]),
-                high=float(row["high"]),
-                low=float(row["low"]),
-                close=float(row["close"]),
-                volume=float(row.get("volume", 0.0))
-            ))
+            bars.append(
+                _rust.Bar(  # type: ignore
+                    timestamp=ts,
+                    open=float(row["open"]),
+                    high=float(row["high"]),
+                    low=float(row["low"]),
+                    close=float(row["close"]),
+                    volume=float(row.get("volume", 0.0)),
+                )
+            )
         return bars
 
     def run(self):
         bars = self._to_bars()
-        engine = _rust.Engine(bars, self._cash, self._commission, self._spread)
+        engine = _rust.Engine(bars, self._cash, self._commission, self._spread)  # type: ignore
         strategy = self._strategy_class()
         return engine.run(_Adapter(strategy))

src/engine.rs

@@ -1,18 +1,18 @@
-use pyo3::prelude::*;
-use crate::types::Bar;
 use crate::broker::Broker;
-use crate::strategy::Strategy;
 use crate::stats::Stats;
+use crate::strategy::Strategy;
+use crate::types::Bar;
+use pyo3::prelude::*;
 
 #[pyclass]
 pub struct Engine {
     pub data: Vec<Bar>,
     pub broker: Broker,
-    pub equity_curve: Vec<f64>
+    pub equity_curve: Vec<f64>,
 }
 
 impl Engine {
-    pub fn run (&mut self, strategy: &mut dyn Strategy) -> Stats {
+    pub fn run(&mut self, strategy: &mut dyn Strategy) -> Stats {
         strategy.init(&self.data);
         for bar in &self.data {
             self.broker.check_sl_tp(bar);
@@ -33,7 +33,7 @@ impl Engine {
         Engine {
             data,
             broker: Broker::new(initial_cash, commission, spread),
-            equity_curve: Vec::new()
+            equity_curve: Vec::new(),
         }
     }
 
@@ -48,18 +48,23 @@ impl Engine {
             }
         }
 
-        let broker_py = Py::new(py, Broker::new(
-            self.broker.initial_cash,
-            self.broker.commission,
-            self.broker.spread,
-        ))?;
+        let broker_py = Py::new(
+            py,
+            Broker::new(
+                self.broker.initial_cash,
+                self.broker.commission,
+                self.broker.spread,
+            ),
+        )?;
 
         for bar in &self.data {
             {
                 let mut b = broker_py.borrow_mut(py);
                 b.check_sl_tp(bar);
             }
-            strategy.bind(py).call_method("next", (bar.clone(), broker_py.clone_ref(py)), None)?;
+            strategy
+                .bind(py)
+                .call_method("next", (bar.clone(), broker_py.clone_ref(py)), None)?;
             let equity = broker_py.borrow(py).equity(bar.close);
             self.equity_curve.push(equity);
         }
@@ -72,4 +77,4 @@ impl Engine {
         let b = broker_py.borrow(py);
         Ok(Stats::compute(&b, &self.equity_curve))
     }
-}
+}