Add own bid and ask quantity methods for OwnOrderBook

Author: cjdsellers

crates/model/src/orderbook/own.rs

@@ -302,6 +302,34 @@ impl OwnOrderBook {
             .collect()
     }
 
+    /// Returns the aggregated own bid quantity at each price level.
+    pub fn bid_quantity(&self) -> IndexMap<Decimal, Decimal> {
+        self.bids()
+            .map(|level| {
+                (
+                    level.price.value.as_decimal(),
+                    level.orders.values().fold(Decimal::ZERO, |total, order| {
+                        total + order.size.as_decimal()
+                    }),
+                )
+            })
+            .collect()
+    }
+
+    /// Returns the aggregated own ask quantity at each price level.
+    pub fn ask_quantity(&self) -> IndexMap<Decimal, Decimal> {
+        self.asks()
+            .map(|level| {
+                (
+                    level.price.value.as_decimal(),
+                    level.orders.values().fold(Decimal::ZERO, |total, order| {
+                        total + order.size.as_decimal()
+                    }),
+                )
+            })
+            .collect()
+    }
+
     /// Return a formatted string representation of the order book.
     #[must_use]
     pub fn pprint(&self, num_levels: usize) -> String {
@@ -587,6 +615,7 @@ impl Ord for OwnBookLevel {
 mod tests {
     use nautilus_core::UnixNanos;
     use rstest::{fixture, rstest};
+    use rust_decimal_macros::dec;
 
     use super::*;
 
@@ -890,4 +919,97 @@ mod tests {
         assert_eq!(ask_orders.len(), 1);
         assert_eq!(ask_orders[0], order2);
     }
+
+    #[rstest]
+    fn test_own_order_book_quantity_empty_levels() {
+        let instrument_id = InstrumentId::from("AAPL.XNAS");
+        let book = OwnOrderBook::new(instrument_id);
+
+        let bid_quantities = book.bid_quantity();
+        let ask_quantities = book.ask_quantity();
+
+        assert!(bid_quantities.is_empty());
+        assert!(ask_quantities.is_empty());
+    }
+
+    #[rstest]
+    fn test_own_order_book_bid_ask_quantity() {
+        let instrument_id = InstrumentId::from("AAPL.XNAS");
+        let mut book = OwnOrderBook::new(instrument_id);
+
+        // Add multiple orders at the same price level (bids)
+        let bid_order1 = OwnBookOrder::new(
+            ClientOrderId::from("O-1"),
+            OrderSideSpecified::Buy,
+            Price::from("100.00"),
+            Quantity::from("10"),
+            OrderType::Limit,
+            TimeInForce::Gtc,
+            OrderStatus::Accepted,
+            UnixNanos::default(),
+            UnixNanos::default(),
+        );
+        let bid_order2 = OwnBookOrder::new(
+            ClientOrderId::from("O-2"),
+            OrderSideSpecified::Buy,
+            Price::from("100.00"),
+            Quantity::from("15"),
+            OrderType::Limit,
+            TimeInForce::Gtc,
+            OrderStatus::Accepted,
+            UnixNanos::default(),
+            UnixNanos::default(),
+        );
+        // Add an order at a different price level (bids)
+        let bid_order3 = OwnBookOrder::new(
+            ClientOrderId::from("O-3"),
+            OrderSideSpecified::Buy,
+            Price::from("99.50"),
+            Quantity::from("20"),
+            OrderType::Limit,
+            TimeInForce::Gtc,
+            OrderStatus::Accepted,
+            UnixNanos::default(),
+            UnixNanos::default(),
+        );
+
+        // Add orders at different price levels (asks)
+        let ask_order1 = OwnBookOrder::new(
+            ClientOrderId::from("O-4"),
+            OrderSideSpecified::Sell,
+            Price::from("101.00"),
+            Quantity::from("12"),
+            OrderType::Limit,
+            TimeInForce::Gtc,
+            OrderStatus::Accepted,
+            UnixNanos::default(),
+            UnixNanos::default(),
+        );
+        let ask_order2 = OwnBookOrder::new(
+            ClientOrderId::from("O-5"),
+            OrderSideSpecified::Sell,
+            Price::from("101.00"),
+            Quantity::from("8"),
+            OrderType::Limit,
+            TimeInForce::Gtc,
+            OrderStatus::Accepted,
+            UnixNanos::default(),
+            UnixNanos::default(),
+        );
+
+        book.add(bid_order1);
+        book.add(bid_order2);
+        book.add(bid_order3);
+        book.add(ask_order1);
+        book.add(ask_order2);
+
+        let bid_quantities = book.bid_quantity();
+        assert_eq!(bid_quantities.len(), 2);
+        assert_eq!(bid_quantities.get(&dec!(100.00)), Some(&dec!(25)));
+        assert_eq!(bid_quantities.get(&dec!(99.50)), Some(&dec!(20)));
+
+        let ask_quantities = book.ask_quantity();
+        assert_eq!(ask_quantities.len(), 1);
+        assert_eq!(ask_quantities.get(&dec!(101.00)), Some(&dec!(20)));
+    }
 }

crates/model/src/python/orderbook/own.rs

@@ -213,6 +213,16 @@ impl OwnOrderBook {
         self.asks_as_map()
     }
 
+    #[pyo3(name = "bid_quantity")]
+    fn py_bid_quantity(&self) -> IndexMap<Decimal, Decimal> {
+        self.bid_quantity()
+    }
+
+    #[pyo3(name = "ask_quantity")]
+    fn py_ask_quantity(&self) -> IndexMap<Decimal, Decimal> {
+        self.ask_quantity()
+    }
+
     #[pyo3(signature = (num_levels=3))]
     #[pyo3(name = "pprint")]
     fn py_pprint(&self, num_levels: usize) -> String {

nautilus_trader/core/nautilus_pyo3.pyi

@@ -3212,6 +3212,8 @@ class OwnOrderBook:
     def clear(self) -> None: ...
     def bids_to_dict(self) -> dict[Decimal, list[OwnBookOrder]]: ...
     def asks_to_dict(self) -> dict[Decimal, list[OwnBookOrder]]: ...
+    def bid_quantity(self) -> dict[Decimal, Decimal]: ...
+    def ask_quantity(self) -> dict[Decimal, Decimal]: ...
     def pprint(self, num_levels: int = 3) -> str: ...
 
 ###################################################################################################

tests/unit_tests/model/test_own_order_book.py

@@ -13,6 +13,8 @@
 #  limitations under the License.
 # -------------------------------------------------------------------------------------------------
 
+from decimal import Decimal
+
 import pytest
 
 from nautilus_trader.core.nautilus_pyo3 import ClientOrderId
@@ -361,3 +363,169 @@ def test_own_order_book_price_change():
     assert len(new_orders) == 1
     assert new_orders[0] == updated
     assert book.ts_last == 11
+
+
+def test_own_order_book_bid_ask_quantity():
+    instrument_id = InstrumentId.from_str("AAPL.XNAS")
+    book = OwnOrderBook(instrument_id)
+
+    # Add multiple orders at the same price level (bids)
+    bid_order1 = OwnBookOrder(
+        client_order_id=ClientOrderId("O-1"),
+        side=OrderSide.BUY,
+        price=Price(100.0, 2),
+        size=Quantity(10.0, 0),
+        order_type=OrderType.LIMIT,
+        time_in_force=TimeInForce.GTC,
+        status=OrderStatus.ACCEPTED,
+        ts_last=0,
+        ts_init=0,
+    )
+    bid_order2 = OwnBookOrder(
+        client_order_id=ClientOrderId("O-2"),
+        side=OrderSide.BUY,
+        price=Price(100.0, 2),
+        size=Quantity(15.0, 0),
+        order_type=OrderType.LIMIT,
+        time_in_force=TimeInForce.GTC,
+        status=OrderStatus.ACCEPTED,
+        ts_last=0,
+        ts_init=0,
+    )
+    # Add an order at a different price level (bids)
+    bid_order3 = OwnBookOrder(
+        client_order_id=ClientOrderId("O-3"),
+        side=OrderSide.BUY,
+        price=Price(99.5, 2),
+        size=Quantity(20.0, 0),
+        order_type=OrderType.LIMIT,
+        time_in_force=TimeInForce.GTC,
+        status=OrderStatus.ACCEPTED,
+        ts_last=0,
+        ts_init=0,
+    )
+
+    # Add orders at different price levels (asks)
+    ask_order1 = OwnBookOrder(
+        client_order_id=ClientOrderId("O-4"),
+        side=OrderSide.SELL,
+        price=Price(101.0, 2),
+        size=Quantity(12.0, 0),
+        order_type=OrderType.LIMIT,
+        time_in_force=TimeInForce.GTC,
+        status=OrderStatus.ACCEPTED,
+        ts_last=0,
+        ts_init=0,
+    )
+    ask_order2 = OwnBookOrder(
+        client_order_id=ClientOrderId("O-5"),
+        side=OrderSide.SELL,
+        price=Price(101.0, 2),
+        size=Quantity(8.0, 0),
+        order_type=OrderType.LIMIT,
+        time_in_force=TimeInForce.GTC,
+        status=OrderStatus.ACCEPTED,
+        ts_last=0,
+        ts_init=0,
+    )
+
+    book.add(bid_order1)
+    book.add(bid_order2)
+    book.add(bid_order3)
+    book.add(ask_order1)
+    book.add(ask_order2)
+
+    bid_quantities = book.bid_quantity()
+    assert len(bid_quantities) == 2
+    assert bid_quantities[Price(100.0, 2).as_decimal()] == Decimal("25")
+    assert bid_quantities[Price(99.5, 2).as_decimal()] == Decimal("20")
+
+    ask_quantities = book.ask_quantity()
+    assert len(ask_quantities) == 1
+    assert ask_quantities[Price(101.0, 2).as_decimal()] == Decimal("20")
+
+
+def test_own_order_book_quantity_empty_levels():
+    instrument_id = InstrumentId.from_str("AAPL.XNAS")
+    book = OwnOrderBook(instrument_id)
+
+    # Test on empty book
+    bid_quantities = book.bid_quantity()
+    ask_quantities = book.ask_quantity()
+
+    assert len(bid_quantities) == 0
+    assert len(ask_quantities) == 0
+
+
+@pytest.mark.parametrize(
+    "orders,expected_bid_quantities,expected_ask_quantities",
+    [
+        # Test case 1: Multiple orders at same price level
+        (
+            [
+                (OrderSide.BUY, 100.0, 10.0),
+                (OrderSide.BUY, 100.0, 15.0),
+                (OrderSide.SELL, 101.0, 20.0),
+            ],
+            {Decimal("100.00"): Decimal("25")},
+            {Decimal("101.00"): Decimal("20")},
+        ),
+        # Test case 2: Multiple price levels
+        (
+            [
+                (OrderSide.BUY, 100.0, 10.0),
+                (OrderSide.BUY, 99.0, 5.0),
+                (OrderSide.SELL, 101.0, 7.0),
+                (OrderSide.SELL, 102.0, 3.0),
+            ],
+            {Decimal("100.00"): Decimal("10"), Decimal("99.00"): Decimal("5")},
+            {Decimal("101.00"): Decimal("7"), Decimal("102.00"): Decimal("3")},
+        ),
+        # Test case 3: Only buy orders
+        (
+            [
+                (OrderSide.BUY, 100.0, 10.0),
+                (OrderSide.BUY, 99.0, 5.0),
+            ],
+            {Decimal("100.00"): Decimal("10"), Decimal("99.00"): Decimal("5")},
+            {},
+        ),
+        # Test case 4: Only sell orders
+        (
+            [
+                (OrderSide.SELL, 101.0, 7.0),
+                (OrderSide.SELL, 102.0, 3.0),
+            ],
+            {},
+            {Decimal("101.00"): Decimal("7"), Decimal("102.00"): Decimal("3")},
+        ),
+    ],
+)
+def test_own_order_book_quantities_parametrized(
+    orders,
+    expected_bid_quantities,
+    expected_ask_quantities,
+):
+    instrument_id = InstrumentId.from_str("AAPL.XNAS")
+    book = OwnOrderBook(instrument_id)
+
+    # Add orders based on the test parameters
+    for i, (side, price, size) in enumerate(orders):
+        order = OwnBookOrder(
+            client_order_id=ClientOrderId(f"O-{i+1}"),
+            side=side,
+            price=Price(price, 2),
+            size=Quantity(size, 0),
+            order_type=OrderType.LIMIT,
+            time_in_force=TimeInForce.GTC,
+            status=OrderStatus.ACCEPTED,
+            ts_last=0,
+            ts_init=0,
+        )
+        book.add(order)
+
+    bid_quantities = book.bid_quantity()
+    ask_quantities = book.ask_quantity()
+
+    assert dict(bid_quantities) == expected_bid_quantities
+    assert dict(ask_quantities) == expected_ask_quantities