nautilus_databento/python/
loader.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
16//! Python bindings for the Databento data loader.
17
18use std::{collections::HashMap, path::PathBuf};
19
20use chrono::NaiveTime;
21use databento::dbn;
22use nautilus_core::python::{IntoPyObjectNautilusExt, to_pyvalue_err};
23use nautilus_model::{
24    data::{
25        Bar, Data, DataFFI, InstrumentStatus, OrderBookDelta, OrderBookDepth10, QuoteTick,
26        TradeTick,
27    },
28    identifiers::{InstrumentId, Symbol, Venue},
29    python::{
30        data::{DATA_FFI_CVEC_CAPSULE_NAME, DataFfiCVec},
31        instruments::instrument_any_to_pyobject,
32    },
33};
34use pyo3::{
35    prelude::*,
36    types::{PyCapsule, PyList},
37};
38use ustr::Ustr;
39
40use crate::{
41    decode::DatabentoDecodeConfig,
42    loader::DatabentoDataLoader,
43    types::{DatabentoImbalance, DatabentoPublisher, DatabentoStatistics, PublisherId},
44};
45
46#[expect(clippy::needless_pass_by_value)]
47#[pymethods]
48#[pyo3_stub_gen::derive::gen_stub_pymethods]
49impl DatabentoDataLoader {
50    /// A Nautilus data loader for Databento Binary Encoding (DBN) format data.
51    ///
52    /// # Supported Schemas
53    ///  - `MBO` -> `OrderBookDelta`
54    ///  - `MBP_1` -> `(QuoteTick, Option<TradeTick>)`
55    ///  - `MBP_10` -> `OrderBookDepth10`
56    ///  - `BBO_1S` -> `QuoteTick`
57    ///  - `BBO_1M` -> `QuoteTick`
58    ///  - `CMBP_1` -> `(QuoteTick, Option<TradeTick>)`
59    ///  - `CBBO_1S` -> `QuoteTick`
60    ///  - `CBBO_1M` -> `QuoteTick`
61    ///  - `TCBBO` -> `(QuoteTick, TradeTick)`
62    ///  - `TBBO` -> `(QuoteTick, TradeTick)`
63    ///  - `TRADES` -> `TradeTick`
64    ///  - `OHLCV_1S` -> `Bar`
65    ///  - `OHLCV_1M` -> `Bar`
66    ///  - `OHLCV_1H` -> `Bar`
67    ///  - `OHLCV_1D` -> `Bar`
68    ///  - `OHLCV_EOD` -> `Bar`
69    ///  - `DEFINITION` -> `Instrument`
70    ///  - `IMBALANCE` -> `DatabentoImbalance`
71    ///  - `STATISTICS` -> `DatabentoStatistics`
72    ///  - `STATUS` -> `InstrumentStatus`
73    ///
74    /// # References
75    ///
76    /// <https://databento.com/docs/schemas-and-data-formats>
77    #[new]
78    #[pyo3(signature = (publishers_filepath=None))]
79    fn py_new(publishers_filepath: Option<PathBuf>) -> PyResult<Self> {
80        Self::new(publishers_filepath).map_err(to_pyvalue_err)
81    }
82
83    /// Load the publishers data from the file at the given `filepath`.
84    ///
85    /// # Errors
86    ///
87    /// Returns an error if the file cannot be read or parsed as JSON.
88    #[pyo3(name = "load_publishers")]
89    fn py_load_publishers(&mut self, publishers_filepath: PathBuf) -> PyResult<()> {
90        self.load_publishers(publishers_filepath)
91            .map_err(to_pyvalue_err)
92    }
93
94    /// Returns the internal Databento publishers currently held by the loader.
95    #[must_use]
96    #[pyo3(name = "get_publishers")]
97    fn py_get_publishers(&self) -> HashMap<u16, DatabentoPublisher> {
98        self.get_publishers()
99            .iter()
100            .map(|(&key, value)| (key, value.clone()))
101            .collect::<HashMap<u16, DatabentoPublisher>>()
102    }
103
104    /// Sets the `venue` to map to the given `dataset`.
105    #[pyo3(name = "set_dataset_for_venue")]
106    fn py_set_dataset_for_venue(&mut self, dataset: String, venue: Venue) {
107        self.set_dataset_for_venue(Ustr::from(&dataset), venue);
108    }
109
110    /// Returns the dataset which matches the given `venue` (if found).
111    #[must_use]
112    #[pyo3(name = "get_dataset_for_venue")]
113    fn py_get_dataset_for_venue(&self, venue: &Venue) -> Option<String> {
114        self.get_dataset_for_venue(venue).map(ToString::to_string)
115    }
116
117    /// Returns the venue which matches the given `publisher_id` (if found).
118    #[must_use]
119    #[pyo3(name = "get_venue_for_publisher")]
120    fn py_get_venue_for_publisher(&self, publisher_id: PublisherId) -> Option<String> {
121        self.get_venue_for_publisher(publisher_id)
122            .map(ToString::to_string)
123    }
124
125    /// Caches a `price_precision` for the given `symbol`.
126    ///
127    /// When market data is read without an explicit `price_precision` argument,
128    /// the loader resolves precision per record from this cache. Definitions
129    /// loaded via `Self.load_instruments` are inserted automatically.
130    #[pyo3(name = "set_price_precision")]
131    fn py_set_price_precision(&mut self, symbol: &str, price_precision: u8) {
132        self.set_price_precision(Symbol::from(symbol), price_precision);
133    }
134
135    /// Returns the cached price precisions keyed by symbol.
136    #[must_use]
137    #[pyo3(name = "get_price_precisions")]
138    fn py_get_price_precisions(&self) -> HashMap<String, u8> {
139        self.get_price_precisions()
140            .iter()
141            .map(|(symbol, precision)| (symbol.to_string(), *precision))
142            .collect()
143    }
144
145    #[pyo3(name = "schema_for_file")]
146    fn py_schema_for_file(&self, filepath: PathBuf) -> PyResult<Option<String>> {
147        self.schema_from_file(&filepath).map_err(to_pyvalue_err)
148    }
149
150    /// Loads all instrument definitions from a DBN file.
151    ///
152    /// When `skip_on_error` is true, instruments that fail to decode are logged
153    /// as warnings and skipped. When false (default), any decode error is propagated.
154    #[pyo3(name = "load_instruments")]
155    #[pyo3(signature = (filepath, use_exchange_as_venue, skip_on_error=false, expiration_overrides=None))]
156    fn py_load_instruments(
157        &mut self,
158        py: Python,
159        filepath: PathBuf,
160        use_exchange_as_venue: bool,
161        skip_on_error: bool,
162        expiration_overrides: Option<HashMap<String, HashMap<String, String>>>,
163    ) -> PyResult<Py<PyAny>> {
164        let decode_config = build_decode_config(expiration_overrides)?;
165        let iter = self
166            .load_instruments(
167                &filepath,
168                use_exchange_as_venue,
169                skip_on_error,
170                decode_config.as_ref(),
171            )
172            .map_err(to_pyvalue_err)?;
173
174        let mut data = Vec::new();
175
176        for instrument in iter {
177            let py_object = instrument_any_to_pyobject(py, instrument)?;
178            data.push(py_object);
179        }
180
181        let list = PyList::new(py, &data).expect("Invalid `ExactSizeIterator`");
182
183        Ok(list.into_py_any_unwrap(py))
184    }
185
186    // Cannot include trades
187    /// Loads order book delta messages from a DBN MBO schema file.
188    ///
189    /// Cannot include trades.
190    #[pyo3(name = "load_order_book_deltas")]
191    #[pyo3(signature = (filepath, instrument_id=None, price_precision=None))]
192    fn py_load_order_book_deltas(
193        &self,
194        filepath: PathBuf,
195        instrument_id: Option<InstrumentId>,
196        price_precision: Option<u8>,
197    ) -> PyResult<Vec<OrderBookDelta>> {
198        self.load_order_book_deltas(&filepath, instrument_id, price_precision)
199            .map_err(to_pyvalue_err)
200    }
201
202    #[pyo3(name = "load_order_book_deltas_as_pycapsule")]
203    #[pyo3(signature = (filepath, instrument_id=None, price_precision=None, include_trades=None))]
204    fn py_load_order_book_deltas_as_pycapsule(
205        &self,
206        py: Python,
207        filepath: PathBuf,
208        instrument_id: Option<InstrumentId>,
209        price_precision: Option<u8>,
210        include_trades: Option<bool>,
211    ) -> PyResult<Py<PyAny>> {
212        let iter = self
213            .read_records::<dbn::MboMsg>(
214                &filepath,
215                instrument_id,
216                price_precision,
217                include_trades.unwrap_or(false),
218                None,
219            )
220            .map_err(to_pyvalue_err)?;
221
222        exhaust_data_iter_to_pycapsule(py, iter).map_err(to_pyvalue_err)
223    }
224
225    /// Loads order book depth10 snapshots from a DBN MBP-10 schema file.
226    #[pyo3(name = "load_order_book_depth10")]
227    #[pyo3(signature = (filepath, instrument_id=None, price_precision=None))]
228    fn py_load_order_book_depth10(
229        &self,
230        filepath: PathBuf,
231        instrument_id: Option<InstrumentId>,
232        price_precision: Option<u8>,
233    ) -> PyResult<Vec<OrderBookDepth10>> {
234        self.load_order_book_depth10(&filepath, instrument_id, price_precision)
235            .map_err(to_pyvalue_err)
236    }
237
238    #[pyo3(name = "load_order_book_depth10_as_pycapsule")]
239    #[pyo3(signature = (filepath, instrument_id=None, price_precision=None))]
240    fn py_load_order_book_depth10_as_pycapsule(
241        &self,
242        py: Python,
243        filepath: PathBuf,
244        instrument_id: Option<InstrumentId>,
245        price_precision: Option<u8>,
246    ) -> PyResult<Py<PyAny>> {
247        let iter = self
248            .read_records::<dbn::Mbp10Msg>(&filepath, instrument_id, price_precision, false, None)
249            .map_err(to_pyvalue_err)?;
250
251        exhaust_data_iter_to_pycapsule(py, iter).map_err(to_pyvalue_err)
252    }
253
254    /// Loads quote tick messages from a DBN MBP-1 or TBBO schema file.
255    #[pyo3(name = "load_quotes")]
256    #[pyo3(signature = (filepath, instrument_id=None, price_precision=None))]
257    fn py_load_quotes(
258        &self,
259        filepath: PathBuf,
260        instrument_id: Option<InstrumentId>,
261        price_precision: Option<u8>,
262    ) -> PyResult<Vec<QuoteTick>> {
263        self.load_quotes(&filepath, instrument_id, price_precision)
264            .map_err(to_pyvalue_err)
265    }
266
267    #[pyo3(name = "load_quotes_as_pycapsule")]
268    #[pyo3(signature = (filepath, instrument_id=None, price_precision=None, include_trades=None))]
269    fn py_load_quotes_as_pycapsule(
270        &self,
271        py: Python,
272        filepath: PathBuf,
273        instrument_id: Option<InstrumentId>,
274        price_precision: Option<u8>,
275        include_trades: Option<bool>,
276    ) -> PyResult<Py<PyAny>> {
277        let iter = self
278            .read_records::<dbn::Mbp1Msg>(
279                &filepath,
280                instrument_id,
281                price_precision,
282                include_trades.unwrap_or(false),
283                None,
284            )
285            .map_err(to_pyvalue_err)?;
286
287        exhaust_data_iter_to_pycapsule(py, iter).map_err(to_pyvalue_err)
288    }
289
290    /// Loads best bid/offer quote messages from a DBN BBO schema file.
291    #[pyo3(name = "load_bbo_quotes")]
292    #[pyo3(signature = (filepath, instrument_id=None, price_precision=None))]
293    fn py_load_bbo_quotes(
294        &self,
295        filepath: PathBuf,
296        instrument_id: Option<InstrumentId>,
297        price_precision: Option<u8>,
298    ) -> PyResult<Vec<QuoteTick>> {
299        self.load_bbo_quotes(&filepath, instrument_id, price_precision)
300            .map_err(to_pyvalue_err)
301    }
302
303    #[pyo3(name = "load_bbo_quotes_as_pycapsule")]
304    #[pyo3(signature = (filepath, instrument_id=None, price_precision=None))]
305    fn py_load_bbo_quotes_as_pycapsule(
306        &self,
307        py: Python,
308        filepath: PathBuf,
309        instrument_id: Option<InstrumentId>,
310        price_precision: Option<u8>,
311    ) -> PyResult<Py<PyAny>> {
312        let iter = self
313            .read_records::<dbn::BboMsg>(&filepath, instrument_id, price_precision, false, None)
314            .map_err(to_pyvalue_err)?;
315
316        exhaust_data_iter_to_pycapsule(py, iter).map_err(to_pyvalue_err)
317    }
318
319    /// Loads consolidated MBP-1 quote messages from a DBN CMBP-1 schema file.
320    #[pyo3(name = "load_cmbp_quotes")]
321    #[pyo3(signature = (filepath, instrument_id=None, price_precision=None))]
322    fn py_load_cmbp_quotes(
323        &self,
324        filepath: PathBuf,
325        instrument_id: Option<InstrumentId>,
326        price_precision: Option<u8>,
327    ) -> PyResult<Vec<QuoteTick>> {
328        self.load_cmbp_quotes(&filepath, instrument_id, price_precision)
329            .map_err(to_pyvalue_err)
330    }
331
332    #[pyo3(name = "load_cmbp_quotes_as_pycapsule")]
333    #[pyo3(signature = (filepath, instrument_id=None, price_precision=None, include_trades=None))]
334    fn py_load_cmbp_quotes_as_pycapsule(
335        &self,
336        py: Python,
337        filepath: PathBuf,
338        instrument_id: Option<InstrumentId>,
339        price_precision: Option<u8>,
340        include_trades: Option<bool>,
341    ) -> PyResult<Py<PyAny>> {
342        let iter = self
343            .read_records::<dbn::Cmbp1Msg>(
344                &filepath,
345                instrument_id,
346                price_precision,
347                include_trades.unwrap_or(false),
348                None,
349            )
350            .map_err(to_pyvalue_err)?;
351
352        exhaust_data_iter_to_pycapsule(py, iter).map_err(to_pyvalue_err)
353    }
354
355    /// Loads consolidated best bid/offer quote messages from a DBN CBBO schema file.
356    #[pyo3(name = "load_cbbo_quotes")]
357    #[pyo3(signature = (filepath, instrument_id=None, price_precision=None))]
358    fn py_load_cbbo_quotes(
359        &self,
360        filepath: PathBuf,
361        instrument_id: Option<InstrumentId>,
362        price_precision: Option<u8>,
363    ) -> PyResult<Vec<QuoteTick>> {
364        self.load_cbbo_quotes(&filepath, instrument_id, price_precision)
365            .map_err(to_pyvalue_err)
366    }
367
368    #[pyo3(name = "load_cbbo_quotes_as_pycapsule")]
369    #[pyo3(signature = (filepath, instrument_id=None, price_precision=None))]
370    fn py_load_cbbo_quotes_as_pycapsule(
371        &self,
372        py: Python,
373        filepath: PathBuf,
374        instrument_id: Option<InstrumentId>,
375        price_precision: Option<u8>,
376    ) -> PyResult<Py<PyAny>> {
377        let iter = self
378            .read_records::<dbn::CbboMsg>(&filepath, instrument_id, price_precision, false, None)
379            .map_err(to_pyvalue_err)?;
380
381        exhaust_data_iter_to_pycapsule(py, iter).map_err(to_pyvalue_err)
382    }
383
384    /// Loads trade messages from a DBN TBBO schema file.
385    #[pyo3(name = "load_tbbo_trades")]
386    #[pyo3(signature = (filepath, instrument_id=None, price_precision=None))]
387    fn py_load_tbbo_trades(
388        &self,
389        filepath: PathBuf,
390        instrument_id: Option<InstrumentId>,
391        price_precision: Option<u8>,
392    ) -> PyResult<Vec<TradeTick>> {
393        self.load_tbbo_trades(&filepath, instrument_id, price_precision)
394            .map_err(to_pyvalue_err)
395    }
396
397    #[pyo3(name = "load_tbbo_trades_as_pycapsule")]
398    #[pyo3(signature = (filepath, instrument_id=None, price_precision=None))]
399    fn py_load_tbbo_trades_as_pycapsule(
400        &self,
401        py: Python,
402        filepath: PathBuf,
403        instrument_id: Option<InstrumentId>,
404        price_precision: Option<u8>,
405    ) -> PyResult<Py<PyAny>> {
406        let iter = self
407            .read_records::<dbn::TbboMsg>(&filepath, instrument_id, price_precision, false, None)
408            .map_err(to_pyvalue_err)?;
409
410        exhaust_data_iter_to_pycapsule(py, iter).map_err(to_pyvalue_err)
411    }
412
413    /// Loads trade messages from a DBN TCBBO schema file.
414    #[pyo3(name = "load_tcbbo_trades")]
415    #[pyo3(signature = (filepath, instrument_id=None, price_precision=None))]
416    fn py_load_tcbbo_trades(
417        &self,
418        filepath: PathBuf,
419        instrument_id: Option<InstrumentId>,
420        price_precision: Option<u8>,
421    ) -> PyResult<Vec<TradeTick>> {
422        self.load_tcbbo_trades(&filepath, instrument_id, price_precision)
423            .map_err(to_pyvalue_err)
424    }
425
426    #[pyo3(name = "load_tcbbo_trades_as_pycapsule")]
427    #[pyo3(signature = (filepath, instrument_id=None, price_precision=None))]
428    fn py_load_tcbbo_trades_as_pycapsule(
429        &self,
430        py: Python,
431        filepath: PathBuf,
432        instrument_id: Option<InstrumentId>,
433        price_precision: Option<u8>,
434    ) -> PyResult<Py<PyAny>> {
435        let iter = self
436            .read_records::<dbn::TcbboMsg>(&filepath, instrument_id, price_precision, false, None)
437            .map_err(to_pyvalue_err)?;
438
439        exhaust_data_iter_to_pycapsule(py, iter).map_err(to_pyvalue_err)
440    }
441
442    /// Loads trade messages from a DBN TRADES schema file.
443    #[pyo3(name = "load_trades")]
444    #[pyo3(signature = (filepath, instrument_id=None, price_precision=None))]
445    fn py_load_trades(
446        &self,
447        filepath: PathBuf,
448        instrument_id: Option<InstrumentId>,
449        price_precision: Option<u8>,
450    ) -> PyResult<Vec<TradeTick>> {
451        self.load_trades(&filepath, instrument_id, price_precision)
452            .map_err(to_pyvalue_err)
453    }
454
455    #[pyo3(name = "load_trades_as_pycapsule")]
456    #[pyo3(signature = (filepath, instrument_id=None, price_precision=None))]
457    fn py_load_trades_as_pycapsule(
458        &self,
459        py: Python,
460        filepath: PathBuf,
461        instrument_id: Option<InstrumentId>,
462        price_precision: Option<u8>,
463    ) -> PyResult<Py<PyAny>> {
464        let iter = self
465            .read_records::<dbn::TradeMsg>(&filepath, instrument_id, price_precision, false, None)
466            .map_err(to_pyvalue_err)?;
467
468        exhaust_data_iter_to_pycapsule(py, iter).map_err(to_pyvalue_err)
469    }
470
471    /// Loads OHLCV bar messages from a DBN OHLCV schema file.
472    #[pyo3(name = "load_bars")]
473    #[pyo3(signature = (filepath, instrument_id=None, price_precision=None, timestamp_on_close=true))]
474    fn py_load_bars(
475        &self,
476        filepath: PathBuf,
477        instrument_id: Option<InstrumentId>,
478        price_precision: Option<u8>,
479        timestamp_on_close: bool,
480    ) -> PyResult<Vec<Bar>> {
481        self.load_bars(
482            &filepath,
483            instrument_id,
484            price_precision,
485            Some(timestamp_on_close),
486        )
487        .map_err(to_pyvalue_err)
488    }
489
490    #[pyo3(name = "load_bars_as_pycapsule")]
491    #[pyo3(signature = (filepath, instrument_id=None, price_precision=None, timestamp_on_close=true))]
492    fn py_load_bars_as_pycapsule(
493        &self,
494        py: Python,
495        filepath: PathBuf,
496        instrument_id: Option<InstrumentId>,
497        price_precision: Option<u8>,
498        timestamp_on_close: bool,
499    ) -> PyResult<Py<PyAny>> {
500        let iter = self
501            .read_records::<dbn::OhlcvMsg>(
502                &filepath,
503                instrument_id,
504                price_precision,
505                false,
506                Some(timestamp_on_close),
507            )
508            .map_err(to_pyvalue_err)?;
509
510        exhaust_data_iter_to_pycapsule(py, iter).map_err(to_pyvalue_err)
511    }
512
513    #[pyo3(name = "load_status")]
514    #[pyo3(signature = (filepath, instrument_id=None))]
515    fn py_load_status(
516        &self,
517        filepath: PathBuf,
518        instrument_id: Option<InstrumentId>,
519    ) -> PyResult<Vec<InstrumentStatus>> {
520        let iter = self
521            .load_status_records::<dbn::StatusMsg>(&filepath, instrument_id)
522            .map_err(to_pyvalue_err)?;
523
524        let mut data = Vec::new();
525
526        for result in iter {
527            match result {
528                Ok(item) => data.push(item),
529                Err(e) => return Err(to_pyvalue_err(e)),
530            }
531        }
532
533        Ok(data)
534    }
535
536    #[pyo3(name = "load_imbalance")]
537    #[pyo3(signature = (filepath, instrument_id=None, price_precision=None))]
538    fn py_load_imbalance(
539        &self,
540        filepath: PathBuf,
541        instrument_id: Option<InstrumentId>,
542        price_precision: Option<u8>,
543    ) -> PyResult<Vec<DatabentoImbalance>> {
544        let iter = self
545            .read_imbalance_records::<dbn::ImbalanceMsg>(&filepath, instrument_id, price_precision)
546            .map_err(to_pyvalue_err)?;
547
548        let mut data = Vec::new();
549
550        for result in iter {
551            match result {
552                Ok(item) => data.push(item),
553                Err(e) => return Err(to_pyvalue_err(e)),
554            }
555        }
556
557        Ok(data)
558    }
559
560    #[pyo3(name = "load_statistics")]
561    #[pyo3(signature = (filepath, instrument_id=None, price_precision=None))]
562    fn py_load_statistics(
563        &self,
564        filepath: PathBuf,
565        instrument_id: Option<InstrumentId>,
566        price_precision: Option<u8>,
567    ) -> PyResult<Vec<DatabentoStatistics>> {
568        let iter = self
569            .read_statistics_records::<dbn::StatMsg>(&filepath, instrument_id, price_precision)
570            .map_err(to_pyvalue_err)?;
571
572        let mut data = Vec::new();
573
574        for result in iter {
575            match result {
576                Ok(item) => data.push(item),
577                Err(e) => return Err(to_pyvalue_err(e)),
578            }
579        }
580
581        Ok(data)
582    }
583}
584
585fn exhaust_data_iter_to_pycapsule(
586    py: Python,
587    iter: impl Iterator<Item = anyhow::Result<(Option<Data>, Option<Data>)>>,
588) -> anyhow::Result<Py<PyAny>> {
589    let mut data = Vec::new();
590
591    for result in iter {
592        match result {
593            Ok((Some(item1), None)) => data.push(item1),
594            Ok((None, Some(item2))) => data.push(item2),
595            Ok((Some(item1), Some(item2))) => {
596                data.push(item1);
597                data.push(item2);
598            }
599            Ok((None, None)) => {}
600            Err(e) => return Err(e),
601        }
602    }
603
604    let ffi_data: Vec<DataFFI> = data
605        .into_iter()
606        .map(DataFFI::try_from)
607        .collect::<Result<Vec<_>, _>>()
608        .map_err(to_pyvalue_err)?;
609    let cvec: DataFfiCVec = ffi_data.into();
610    // No destructor: Python must call drop_cvec_pycapsule to take ownership and free.
611    let capsule = PyCapsule::new_with_value_and_destructor::<DataFfiCVec, _>(
612        py,
613        cvec,
614        DATA_FFI_CVEC_CAPSULE_NAME,
615        |_, _| {},
616    )?;
617
618    // TODO: Improve error domain. Replace anyhow errors with nautilus
619    // errors to unify pyo3 and anyhow errors.
620    Ok(capsule.into_py_any_unwrap(py))
621}
622
623// Returns `None` when no overrides are supplied, so the loader applies its built-in defaults
624fn build_decode_config(
625    expiration_overrides: Option<HashMap<String, HashMap<String, String>>>,
626) -> PyResult<Option<DatabentoDecodeConfig>> {
627    expiration_overrides
628        .map(|overrides| decode_config_from_overrides(overrides).map_err(to_pyvalue_err))
629        .transpose()
630}
631
632// Builds a decode config from a dataset -> (underlying -> wall-clock time) mapping. The reserved
633// underlying key "default" sets a dataset's default time; other keys are per-underlying overrides.
634fn decode_config_from_overrides(
635    overrides: HashMap<String, HashMap<String, String>>,
636) -> Result<DatabentoDecodeConfig, String> {
637    let mut config = DatabentoDecodeConfig::default();
638
639    for (dataset_name, times) in overrides {
640        let dataset = dataset_name
641            .parse::<dbn::Dataset>()
642            .map_err(|_| format!("Unknown dataset '{dataset_name}'"))?;
643        let rule = config
644            .option_expiration
645            .get_mut(&dataset)
646            .ok_or_else(|| format!("No expiration correction rule for dataset '{dataset_name}'"))?;
647
648        for (underlying, value) in times {
649            let time = parse_expiration_time(&value)?;
650            if underlying == "default" {
651                rule.default_time = time;
652            } else {
653                rule.overrides.insert(Ustr::from(&underlying), time);
654            }
655        }
656    }
657
658    Ok(config)
659}
660
661// Parses a `HH:MM` or `HH:MM:SS` wall-clock time
662fn parse_expiration_time(value: &str) -> Result<NaiveTime, String> {
663    NaiveTime::parse_from_str(value, "%H:%M:%S")
664        .or_else(|_| NaiveTime::parse_from_str(value, "%H:%M"))
665        .map_err(|_| format!("Invalid expiration time '{value}', expected 'HH:MM' or 'HH:MM:SS'"))
666}
667
668#[cfg(test)]
669mod tests {
670    use std::collections::HashMap;
671
672    use chrono::NaiveTime;
673    use databento::dbn;
674    use rstest::rstest;
675    use ustr::Ustr;
676
677    use super::decode_config_from_overrides;
678
679    fn overrides(
680        dataset: &str,
681        entries: &[(&str, &str)],
682    ) -> HashMap<String, HashMap<String, String>> {
683        let inner = entries
684            .iter()
685            .map(|(k, v)| ((*k).to_string(), (*v).to_string()))
686            .collect();
687        HashMap::from([(dataset.to_string(), inner)])
688    }
689
690    #[rstest]
691    fn test_default_key_sets_dataset_default_time() {
692        let config =
693            decode_config_from_overrides(overrides("OPRA.PILLAR", &[("default", "15:30")]))
694                .unwrap();
695        let rule = config
696            .option_expiration
697            .get(&dbn::Dataset::OpraPillar)
698            .unwrap();
699        assert_eq!(
700            rule.default_time,
701            NaiveTime::from_hms_opt(15, 30, 0).unwrap()
702        );
703        assert!(rule.overrides.is_empty());
704    }
705
706    #[rstest]
707    fn test_underlying_key_sets_override_and_keeps_default() {
708        let config =
709            decode_config_from_overrides(overrides("OPRA.PILLAR", &[("SPX", "09:30:00")])).unwrap();
710        let rule = config
711            .option_expiration
712            .get(&dbn::Dataset::OpraPillar)
713            .unwrap();
714        assert_eq!(
715            rule.overrides.get(&Ustr::from("SPX")).copied(),
716            Some(NaiveTime::from_hms_opt(9, 30, 0).unwrap())
717        );
718        assert_eq!(
719            rule.default_time,
720            NaiveTime::from_hms_opt(16, 0, 0).unwrap()
721        );
722    }
723
724    #[rstest]
725    fn test_unknown_dataset_errors() {
726        let err = decode_config_from_overrides(overrides("NOT.ADATASET", &[("default", "16:00")]))
727            .unwrap_err();
728        assert!(err.contains("Unknown dataset"), "was: {err}");
729    }
730
731    #[rstest]
732    fn test_dataset_without_rule_errors() {
733        // GLBX is a valid dataset but has no built-in expiration rule, so it cannot be tuned
734        let err = decode_config_from_overrides(overrides("GLBX.MDP3", &[("default", "16:00")]))
735            .unwrap_err();
736        assert!(err.contains("No expiration correction rule"), "was: {err}");
737    }
738
739    #[rstest]
740    fn test_invalid_time_errors() {
741        let err = decode_config_from_overrides(overrides("OPRA.PILLAR", &[("default", "nope")]))
742            .unwrap_err();
743        assert!(err.contains("Invalid expiration time"), "was: {err}");
744    }
745}
nautilus_databento/python/loader.rs

nautilus_databento/python/
loader.rs
