nautilus_tardis/python/
machine.rs

1// -------------------------------------------------------------------------------------------------
2//  Copyright (C) 2015-2026 Nautech Systems Pty Ltd. All rights reserved.
3//  https://nautechsystems.io
4//
5//  Licensed under the GNU Lesser General Public License Version 3.0 (the "License");
6//  You may not use this file except in compliance with the License.
7//  You may obtain a copy of the License at https://www.gnu.org/licenses/lgpl-3.0.en.html
8//
9//  Unless required by applicable law or agreed to in writing, software
10//  distributed under the License is distributed on an "AS IS" BASIS,
11//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12//  See the License for the specific language governing permissions and
13//  limitations under the License.
14// -------------------------------------------------------------------------------------------------
15
16use std::{path::Path, sync::Arc};
17
18use ahash::AHashMap;
19use futures_util::{Stream, StreamExt, pin_mut};
20use nautilus_core::python::{IntoPyObjectNautilusExt, call_python, to_pyruntime_err};
21use nautilus_model::{
22    data::{Bar, Data, funding::FundingRateUpdate},
23    identifiers::InstrumentId,
24    python::data::data_to_pycapsule,
25};
26use pyo3::{prelude::*, types::PyList};
27
28use crate::{
29    config::BookSnapshotOutput,
30    machine::{
31        Error,
32        client::{TardisMachineClient, determine_instrument_info},
33        message::WsMessage,
34        parse::{
35            parse_tardis_ws_message, parse_tardis_ws_message_data,
36            parse_tardis_ws_message_funding_rate,
37        },
38        replay_normalized, stream_normalized,
39        types::{
40            ReplayNormalizedRequestOptions, StreamNormalizedRequestOptions, TardisInstrumentKey,
41            TardisInstrumentMiniInfo,
42        },
43    },
44    replay::run_tardis_machine_replay_from_config,
45};
46
47#[pymethods]
48#[pyo3_stub_gen::derive::gen_stub_pymethods]
49impl ReplayNormalizedRequestOptions {
50    #[staticmethod]
51    #[pyo3(name = "from_json")]
52    fn py_from_json(#[gen_stub(override_type(type_repr = "bytes"))] data: &[u8]) -> Self {
53        serde_json::from_slice(data).expect("Failed to parse JSON")
54    }
55
56    #[pyo3(name = "from_json_array")]
57    #[staticmethod]
58    fn py_from_json_array(
59        #[gen_stub(override_type(type_repr = "bytes"))] data: &[u8],
60    ) -> Vec<Self> {
61        serde_json::from_slice(data).expect("Failed to parse JSON array")
62    }
63}
64
65#[pymethods]
66#[pyo3_stub_gen::derive::gen_stub_pymethods]
67impl StreamNormalizedRequestOptions {
68    #[staticmethod]
69    #[pyo3(name = "from_json")]
70    fn py_from_json(#[gen_stub(override_type(type_repr = "bytes"))] data: &[u8]) -> Self {
71        serde_json::from_slice(data).expect("Failed to parse JSON")
72    }
73
74    #[pyo3(name = "from_json_array")]
75    #[staticmethod]
76    fn py_from_json_array(
77        #[gen_stub(override_type(type_repr = "bytes"))] data: &[u8],
78    ) -> Vec<Self> {
79        serde_json::from_slice(data).expect("Failed to parse JSON array")
80    }
81}
82
83#[pymethods]
84#[pyo3_stub_gen::derive::gen_stub_pymethods]
85impl TardisMachineClient {
86    /// Provides a client for connecting to a [Tardis Machine Server](https://docs.tardis.dev/api/tardis-machine).
87    #[new]
88    #[pyo3(signature = (
89        base_url = None,
90        normalize_symbols = true,
91        book_snapshot_output = "deltas",
92        extract_bbo_as_quotes = false,
93    ))]
94    fn py_new(
95        base_url: Option<&str>,
96        normalize_symbols: bool,
97        book_snapshot_output: &str,
98        extract_bbo_as_quotes: bool,
99    ) -> PyResult<Self> {
100        let output = match book_snapshot_output {
101            "depth10" => BookSnapshotOutput::Depth10,
102            "deltas" => BookSnapshotOutput::Deltas,
103            _ => {
104                return Err(to_pyruntime_err(anyhow::anyhow!(
105                    "Invalid book_snapshot_output: '{book_snapshot_output}'. Expected 'depth10' or 'deltas'"
106                )));
107            }
108        };
109        let mut client =
110            Self::new(base_url, normalize_symbols, output).map_err(to_pyruntime_err)?;
111        client.extract_bbo_as_quotes = extract_bbo_as_quotes;
112        Ok(client)
113    }
114
115    /// Returns `true` if `close()` has been called.
116    ///
117    /// This checks that both replay and stream signals have been set,
118    /// which only occurs when `close()` is explicitly called.
119    #[pyo3(name = "is_closed")]
120    #[must_use]
121    pub fn py_is_closed(&self) -> bool {
122        self.is_closed()
123    }
124
125    #[pyo3(name = "close")]
126    fn py_close(&mut self) {
127        self.close();
128    }
129
130    /// Connects to the Tardis Machine replay WebSocket and yields parsed `Data` items.
131    #[pyo3(name = "replay")]
132    fn py_replay<'py>(
133        &self,
134        instruments: Vec<TardisInstrumentMiniInfo>,
135        options: Vec<ReplayNormalizedRequestOptions>,
136        callback: Py<PyAny>,
137        py: Python<'py>,
138    ) -> PyResult<Bound<'py, PyAny>> {
139        let map = if instruments.is_empty() {
140            self.instruments.clone()
141        } else {
142            let mut instrument_map: AHashMap<TardisInstrumentKey, Arc<TardisInstrumentMiniInfo>> =
143                AHashMap::new();
144
145            for inst in instruments {
146                let key = inst.as_tardis_instrument_key();
147                instrument_map.insert(key, Arc::new(inst.clone()));
148            }
149            instrument_map
150        };
151
152        let base_url = self.base_url.clone();
153        let replay_signal = self.replay_signal.clone();
154        let book_snapshot_output = self.book_snapshot_output.clone();
155        let extract_bbo_as_quotes = self.extract_bbo_as_quotes;
156
157        pyo3_async_runtimes::tokio::future_into_py(py, async move {
158            let stream = replay_normalized(&base_url, options, replay_signal)
159                .await
160                .map_err(to_pyruntime_err)?;
161
162            // We use Box::pin to heap-allocate the stream and ensure it implements
163            // Unpin for safe async handling across lifetimes.
164            handle_python_stream(
165                Box::pin(stream),
166                callback,
167                None,
168                Some(map),
169                book_snapshot_output,
170                extract_bbo_as_quotes,
171            )
172            .await;
173            Ok(())
174        })
175    }
176
177    #[pyo3(name = "replay_bars")]
178    fn py_replay_bars<'py>(
179        &self,
180        instruments: Vec<TardisInstrumentMiniInfo>,
181        options: Vec<ReplayNormalizedRequestOptions>,
182        py: Python<'py>,
183    ) -> PyResult<Bound<'py, PyAny>> {
184        let map = if instruments.is_empty() {
185            self.instruments.clone()
186        } else {
187            instruments
188                .into_iter()
189                .map(|inst| (inst.as_tardis_instrument_key(), Arc::new(inst)))
190                .collect()
191        };
192
193        let base_url = self.base_url.clone();
194        let replay_signal = self.replay_signal.clone();
195        let book_snapshot_output = self.book_snapshot_output.clone();
196
197        pyo3_async_runtimes::tokio::future_into_py(py, async move {
198            let stream = replay_normalized(&base_url, options, replay_signal)
199                .await
200                .map_err(to_pyruntime_err)?;
201
202            // We use Box::pin to heap-allocate the stream and ensure it implements
203            // Unpin for safe async handling across lifetimes.
204            pin_mut!(stream);
205
206            let mut bars: Vec<Bar> = Vec::new();
207
208            while let Some(result) = stream.next().await {
209                match result {
210                    Ok(msg) => {
211                        if let Some(Data::Bar(bar)) = determine_instrument_info(&msg, &map)
212                            .and_then(|info| {
213                                parse_tardis_ws_message(msg, &info, &book_snapshot_output)
214                            })
215                        {
216                            bars.push(bar);
217                        }
218                    }
219                    Err(e) => {
220                        log::error!("Error in WebSocket stream: {e:?}");
221                        break;
222                    }
223                }
224            }
225
226            Python::attach(|py| {
227                let pylist =
228                    PyList::new(py, bars.into_iter().map(|bar| bar.into_py_any_unwrap(py)))
229                        .expect("Invalid `ExactSizeIterator`");
230                Ok(pylist.into_py_any_unwrap(py))
231            })
232        })
233    }
234
235    /// Connects to the Tardis Machine stream WebSocket for a single instrument and yields parsed `Data` items.
236    #[pyo3(name = "stream")]
237    fn py_stream<'py>(
238        &self,
239        instruments: Vec<TardisInstrumentMiniInfo>,
240        options: Vec<StreamNormalizedRequestOptions>,
241        callback: Py<PyAny>,
242        py: Python<'py>,
243    ) -> PyResult<Bound<'py, PyAny>> {
244        let mut instrument_map: AHashMap<TardisInstrumentKey, Arc<TardisInstrumentMiniInfo>> =
245            AHashMap::new();
246
247        for inst in instruments {
248            let key = inst.as_tardis_instrument_key();
249            instrument_map.insert(key, Arc::new(inst.clone()));
250        }
251
252        let base_url = self.base_url.clone();
253        let replay_signal = self.replay_signal.clone();
254        let book_snapshot_output = self.book_snapshot_output.clone();
255        let extract_bbo_as_quotes = self.extract_bbo_as_quotes;
256
257        pyo3_async_runtimes::tokio::future_into_py(py, async move {
258            let stream = stream_normalized(&base_url, options, replay_signal)
259                .await
260                .map_err(to_pyruntime_err)?;
261
262            // We use Box::pin to heap-allocate the stream and ensure it implements
263            // Unpin for safe async handling across lifetimes.
264            handle_python_stream(
265                Box::pin(stream),
266                callback,
267                None,
268                Some(instrument_map),
269                book_snapshot_output,
270                extract_bbo_as_quotes,
271            )
272            .await;
273            Ok(())
274        })
275    }
276}
277
278/// Run the Tardis Machine replay as an async Python future.
279///
280/// # Errors
281///
282/// Returns a `PyErr` if reading the config file or replay execution fails.
283#[pyfunction]
284#[pyo3_stub_gen::derive::gen_stub_pyfunction(module = "nautilus_trader.adapters.tardis")]
285#[pyo3(name = "run_tardis_machine_replay")]
286#[pyo3(signature = (config_filepath))]
287pub fn py_run_tardis_machine_replay(
288    py: Python<'_>,
289    config_filepath: String,
290) -> PyResult<Bound<'_, PyAny>> {
291    nautilus_common::logging::ensure_logging_initialized();
292
293    pyo3_async_runtimes::tokio::future_into_py(py, async move {
294        let config_filepath = Path::new(&config_filepath);
295        run_tardis_machine_replay_from_config(config_filepath)
296            .await
297            .map_err(to_pyruntime_err)?;
298        Ok(())
299    })
300}
301
302async fn handle_python_stream<S>(
303    stream: S,
304    callback: Py<PyAny>,
305    instrument: Option<Arc<TardisInstrumentMiniInfo>>,
306    instrument_map: Option<AHashMap<TardisInstrumentKey, Arc<TardisInstrumentMiniInfo>>>,
307    book_snapshot_output: BookSnapshotOutput,
308    extract_bbo_as_quotes: bool,
309) where
310    S: Stream<Item = Result<WsMessage, Error>> + Unpin,
311{
312    pin_mut!(stream);
313
314    // Cache for funding rates to avoid duplicate emissions
315    let mut funding_rate_cache: AHashMap<InstrumentId, FundingRateUpdate> = AHashMap::new();
316
317    while let Some(result) = stream.next().await {
318        match result {
319            Ok(msg) => {
320                let info = instrument.clone().or_else(|| {
321                    instrument_map
322                        .as_ref()
323                        .and_then(|map| determine_instrument_info(&msg, map))
324                });
325
326                if let Some(info) = info.clone() {
327                    let data = parse_tardis_ws_message_data(
328                        msg.clone(),
329                        &info,
330                        &book_snapshot_output,
331                        extract_bbo_as_quotes,
332                    );
333
334                    if !data.is_empty() {
335                        Python::attach(|py| {
336                            for data in data {
337                                let py_obj = data_to_pycapsule(py, data);
338                                call_python(py, &callback, py_obj);
339                            }
340                        });
341                    } else if let Some(funding_rate) =
342                        parse_tardis_ws_message_funding_rate(msg, &info)
343                    {
344                        // Check if we should emit this funding rate
345                        let should_emit = if let Some(cached_rate) =
346                            funding_rate_cache.get(&funding_rate.instrument_id)
347                        {
348                            // Only emit if changed (uses custom PartialEq comparing rate and next_funding_ns)
349                            if cached_rate == &funding_rate {
350                                false // Skip unchanged rate
351                            } else {
352                                funding_rate_cache.insert(funding_rate.instrument_id, funding_rate);
353                                true
354                            }
355                        } else {
356                            // First time seeing this instrument, cache and emit
357                            funding_rate_cache.insert(funding_rate.instrument_id, funding_rate);
358                            true
359                        };
360
361                        if should_emit {
362                            Python::attach(|py| {
363                                let py_obj = funding_rate.into_py_any_unwrap(py);
364                                call_python(py, &callback, py_obj);
365                            });
366                        }
367                    }
368                }
369            }
370            Err(e) => {
371                log::error!("Error in WebSocket stream: {e:?}");
372                break;
373            }
374        }
375    }
376}
nautilus_tardis/python/machine.rs

nautilus_tardis/python/
machine.rs
