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

nautilus_tardis/python/
machine.rs
