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

nautilus_tardis/python/
machine.rs
