core: fully restructured code to support pyo3

src/broker.rs

@@ -1,6 +1,8 @@
 
+use pyo3::prelude::*;
 use crate::types::{Position, Trade, Bar};
 
+#[pyclass]
 pub struct Broker{
     pub cash: f64,
     pub initial_cash: f64,
@@ -10,115 +12,8 @@ pub struct Broker{
     pub spread: f64
 }
 
+// Rust-internal only, not exposed to Python
 impl Broker {
-    pub fn new(initial_cash: f64, commission: f64, spread: f64) -> Self {     // does not need &mut because it initialises something new 
-        Broker {
-            cash: initial_cash,
-            initial_cash,
-            positions: Vec::new(),
-            trade_history: Vec::new(),
-            commission,
-            spread
-        }
-    }
-
-    pub fn buy(&mut self, price: f64, lot_size: f64, timestamp: i64, stop_loss: Option<f64>, take_profit: Option<f64>) {      // needs to modify the broker with new position. (.push works with the Vec::)
-        let fill_price = price + self.spread; // buy at ask 
-        self.cash -= self.commission * lot_size; // pay commission
-        self.positions.push(Position {
-            id: self.positions.len() as u64,
-            entry_price: fill_price,
-            lot_size,
-            is_long: true,
-            entry_timestamp: timestamp,
-            stop_loss,
-            take_profit
-        });
-    }
-
-    pub fn sell(&mut self, price: f64, lot_size: f64, timestamp: i64, stop_loss: Option<f64>, take_profit: Option<f64>) {  
-        let fill_price = price - self.spread; // buy at ask 
-        self.cash -= self.commission * lot_size; // pay commission   // needs to modify the broker with new position. (.push works with the Vec::)
-        self.positions.push(Position {
-            id : self.positions.len() as u64,
-            entry_price: fill_price,
-            lot_size,
-            is_long: false,
-            entry_timestamp: timestamp,
-            stop_loss,
-            take_profit
-        });
-    }
-
-    pub fn close_position (&mut self, id: u64, price:f64, timestamp:i64) {
-        if let Some(index) = self.positions.iter().position(|p| p.id == id) {
-            let position = self.positions.remove(index);
-
-            
-            let close_price = if position.is_long {
-                price - self.spread // sell at the bif whilst closing
-            } else {
-                price + self.spread
-            };
-            
-            let pnl = if position.is_long {
-                (close_price - position.entry_price) * position.lot_size
-            } else {
-                (position.entry_price - close_price) * position.lot_size
-            };
-
-            self.cash += pnl;
-            self.trade_history.push(Trade {
-                entry_price: position.entry_price,
-                lot_size: position.lot_size,
-                is_long: position.is_long,
-                pnl,
-                entry_timestamp: position.entry_timestamp,
-                exit_timestamp: timestamp,
-                exit_price: close_price
-            });            
-        }
-    }
-
-    pub fn close_all (&mut self, price:f64, timestamp: i64) {       // Instead of .push it uses drain to calculate the close all positions
-        for position in self.positions.drain(..) {
-
-            let close_price = if position.is_long {
-                price - self.spread // sell at the bif whilst closing
-            } else {
-                price + self.spread
-            };
-
-            let pnl = if position.is_long {
-                (close_price - position.entry_price) * position.lot_size
-            } else {
-                (position.entry_price - close_price) * position.lot_size
-            };
-            self.cash += pnl;
-            self.trade_history.push(Trade {
-                entry_price: position.entry_price,
-                lot_size: position.lot_size,
-                is_long: position.is_long,
-                pnl,
-                entry_timestamp: position.entry_timestamp,
-                exit_timestamp: timestamp,
-                exit_price: close_price
-            });
-        }
-    }
-
-    pub fn equity(&self, current_price: f64) -> f64 { // computes unrealised positions from the opened positions
-        let unrealized: f64 = self.positions.iter().map(|p| {
-            if p.is_long {
-                (current_price - p.entry_price) * p.lot_size
-            } else {
-                (p.entry_price - current_price) * p.lot_size
-            }
-        }).sum();
-
-        self.cash + unrealized
-    }
-    
     pub fn check_sl_tp(&mut self, bar: &Bar) {
         let mut i = 0;
         while i < self.positions.len() {
@@ -164,10 +59,117 @@ impl Broker {
                 });
             } else {
                 i+=1;
-            } 
-
+            }
         }
-
     }
-    
+}
+
+#[pymethods]
+impl Broker {
+    #[new]
+    pub fn new(initial_cash: f64, commission: f64, spread: f64) -> Self {
+        Broker {
+            cash: initial_cash,
+            initial_cash,
+            positions: Vec::new(),
+            trade_history: Vec::new(),
+            commission,
+            spread
+        }
+    }
+
+    pub fn buy(&mut self, price: f64, lot_size: f64, timestamp: i64, stop_loss: Option<f64>, take_profit: Option<f64>) {
+        let fill_price = price + self.spread;
+        self.cash -= self.commission * lot_size;
+        self.positions.push(Position {
+            id: self.positions.len() as u64,
+            entry_price: fill_price,
+            lot_size,
+            is_long: true,
+            entry_timestamp: timestamp,
+            stop_loss,
+            take_profit
+        });
+    }
+
+    pub fn sell(&mut self, price: f64, lot_size: f64, timestamp: i64, stop_loss: Option<f64>, take_profit: Option<f64>) {
+        let fill_price = price - self.spread;
+        self.cash -= self.commission * lot_size;
+        self.positions.push(Position {
+            id: self.positions.len() as u64,
+            entry_price: fill_price,
+            lot_size,
+            is_long: false,
+            entry_timestamp: timestamp,
+            stop_loss,
+            take_profit
+        });
+    }
+
+    pub fn close_position(&mut self, id: u64, price: f64, timestamp: i64) {
+        if let Some(index) = self.positions.iter().position(|p| p.id == id) {
+            let position = self.positions.remove(index);
+
+            let close_price = if position.is_long {
+                price - self.spread
+            } else {
+                price + self.spread
+            };
+
+            let pnl = if position.is_long {
+                (close_price - position.entry_price) * position.lot_size
+            } else {
+                (position.entry_price - close_price) * position.lot_size
+            };
+
+            self.cash += pnl;
+            self.trade_history.push(Trade {
+                entry_price: position.entry_price,
+                lot_size: position.lot_size,
+                is_long: position.is_long,
+                pnl,
+                entry_timestamp: position.entry_timestamp,
+                exit_timestamp: timestamp,
+                exit_price: close_price
+            });
+        }
+    }
+
+    pub fn close_all(&mut self, price: f64, timestamp: i64) {
+        for position in self.positions.drain(..) {
+            let close_price = if position.is_long {
+                price - self.spread
+            } else {
+                price + self.spread
+            };
+
+            let pnl = if position.is_long {
+                (close_price - position.entry_price) * position.lot_size
+            } else {
+                (position.entry_price - close_price) * position.lot_size
+            };
+            self.cash += pnl;
+            self.trade_history.push(Trade {
+                entry_price: position.entry_price,
+                lot_size: position.lot_size,
+                is_long: position.is_long,
+                pnl,
+                entry_timestamp: position.entry_timestamp,
+                exit_timestamp: timestamp,
+                exit_price: close_price
+            });
+        }
+    }
+
+    pub fn equity(&self, current_price: f64) -> f64 {
+        let unrealized: f64 = self.positions.iter().map(|p| {
+            if p.is_long {
+                (current_price - p.entry_price) * p.lot_size
+            } else {
+                (p.entry_price - current_price) * p.lot_size
+            }
+        }).sum();
+
+        self.cash + unrealized
+    }
 }

src/engine.rs

@@ -1,8 +1,10 @@
+use pyo3::prelude::*;
 use crate::types::Bar;
 use crate::broker::Broker;
 use crate::strategy::Strategy;
 use crate::stats::Stats;
 
+#[pyclass]
 pub struct Engine {
     pub data: Vec<Bar>,
     pub broker: Broker,
@@ -10,14 +12,6 @@ pub struct Engine {
 }
 
 impl Engine {
-    pub fn new(data: Vec<Bar>, initial_cash: f64, commission: f64, spread: f64) -> Self {
-        Engine {
-            data,
-            broker : Broker::new(initial_cash, commission, spread), // it takes initial cash and not Broker as Engine is responsible for broker not the user 
-            equity_curve: Vec::new()
-        }
-    }
-
     pub fn run (&mut self, strategy: &mut dyn Strategy) -> Stats {
         strategy.init(&self.data);
         for bar in &self.data {
@@ -30,4 +24,52 @@ impl Engine {
         }
         Stats::compute(&self.broker, &self.equity_curve)
     }
+}
+
+#[pymethods]
+impl Engine {
+    #[new]
+    pub fn new(data: Vec<Bar>, initial_cash: f64, commission: f64, spread: f64) -> Self {
+        Engine {
+            data,
+            broker: Broker::new(initial_cash, commission, spread),
+            equity_curve: Vec::new()
+        }
+    }
+
+    #[pyo3(name = "run")]
+    pub fn run_py(&mut self, py: Python<'_>, strategy: Py<PyAny>) -> PyResult<Stats> {
+        self.equity_curve.clear();
+
+        let init_result = strategy.bind(py).call_method1("init", (self.data.clone(),));
+        if let Err(e) = init_result {
+            if !e.is_instance_of::<pyo3::exceptions::PyAttributeError>(py) {
+                return Err(e);
+            }
+        }
+
+        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)?;
+            let equity = broker_py.borrow(py).equity(bar.close);
+            self.equity_curve.push(equity);
+        }
+
+        if let Some(last_bar) = self.data.last() {
+            let mut b = broker_py.borrow_mut(py);
+            b.close_all(last_bar.close, last_bar.timestamp);
+        }
+
+        let b = broker_py.borrow(py);
+        Ok(Stats::compute(&b, &self.equity_curve))
+    }
 }

src/lib.rs

@@ -11,5 +11,7 @@ use pyo3::prelude::*;
 fn backtestingfx(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_class::<types::Bar>()?;
     m.add_class::<stats::Stats>()?;
+    m.add_class::<broker::Broker>()?;
+    m.add_class::<engine::Engine>()?;
     Ok(())
 }

src/types.rs

@@ -18,6 +18,14 @@ pub struct Bar {
     pub volume: f64,
 }
 
+#[pymethods]
+impl Bar {
+    #[new]
+    pub fn new(timestamp: i64, open: f64, high: f64, low: f64, close: f64, volume: f64) -> Self {
+        Bar { timestamp, open, high, low, close, volume }
+    }
+}
+
 #[derive(Debug, Clone)]
 pub struct Position {
     // this is for the trading position