nautilus_persistence/python/wranglers/
depth.rs1use std::{collections::HashMap, io::Cursor, str::FromStr};
17
18use datafusion::arrow::ipc::reader::StreamReader;
19use nautilus_core::python::to_pyvalue_err;
20use nautilus_model::{data::OrderBookDepth10, identifiers::InstrumentId};
21use nautilus_serialization::arrow::DecodeFromRecordBatch;
22use pyo3::prelude::*;
23
24#[pyclass]
25#[pyo3_stub_gen::derive::gen_stub_pyclass(module = "nautilus_trader.persistence")]
26pub struct OrderBookDepth10DataWrangler {
27 instrument_id: InstrumentId,
28 price_precision: u8,
29 size_precision: u8,
30 metadata: HashMap<String, String>,
31}
32
33#[pymethods]
34#[pyo3_stub_gen::derive::gen_stub_pymethods]
35impl OrderBookDepth10DataWrangler {
36 #[new]
37 fn py_new(instrument_id: &str, price_precision: u8, size_precision: u8) -> PyResult<Self> {
38 let instrument_id = InstrumentId::from_str(instrument_id).map_err(to_pyvalue_err)?;
39 let metadata =
40 OrderBookDepth10::get_metadata(&instrument_id, price_precision, size_precision);
41
42 Ok(Self {
43 instrument_id,
44 price_precision,
45 size_precision,
46 metadata,
47 })
48 }
49
50 #[getter]
51 fn instrument_id(&self) -> String {
52 self.instrument_id.to_string()
53 }
54
55 #[getter]
56 const fn price_precision(&self) -> u8 {
57 self.price_precision
58 }
59
60 #[getter]
61 const fn size_precision(&self) -> u8 {
62 self.size_precision
63 }
64
65 fn process_record_batch_bytes(
66 &self,
67 #[gen_stub(override_type(type_repr = "bytes"))] data: &[u8],
68 ) -> PyResult<Vec<OrderBookDepth10>> {
69 let cursor = Cursor::new(data);
71 let reader = match StreamReader::try_new(cursor, None) {
72 Ok(reader) => reader,
73 Err(e) => return Err(to_pyvalue_err(e)),
74 };
75
76 let mut depths = Vec::new();
77
78 for maybe_batch in reader {
80 let record_batch = match maybe_batch {
81 Ok(record_batch) => record_batch,
82 Err(e) => return Err(to_pyvalue_err(e)),
83 };
84
85 let batch_depths = OrderBookDepth10::decode_batch(&self.metadata, record_batch)
86 .map_err(to_pyvalue_err)?;
87 depths.extend(batch_depths);
88 }
89
90 Ok(depths)
91 }
92}