nautilus_databento/python/

loader.rs

// -------------------------------------------------------------------------------------------------
//  Copyright (C) 2015-2026 Nautech Systems Pty Ltd. All rights reserved.
//  https://nautechsystems.io
//
//  Licensed under the GNU Lesser General Public License Version 3.0 (the "License");
//  You may not use this file except in compliance with the License.
//  You may obtain a copy of the License at https://www.gnu.org/licenses/lgpl-3.0.en.html
//
//  Unless required by applicable law or agreed to in writing, software
//  distributed under the License is distributed on an "AS IS" BASIS,
//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//  See the License for the specific language governing permissions and
//  limitations under the License.
// -------------------------------------------------------------------------------------------------

//! Python bindings for the Databento data loader.

use std::{collections::HashMap, path::PathBuf};

use databento::dbn;
use jiff::civil::Time;
use nautilus_core::python::{IntoPyObjectNautilusExt, to_pyvalue_err};
use nautilus_model::{
    data::{Bar, InstrumentStatus, OrderBookDelta, OrderBookDepth10, QuoteTick, TradeTick},
    identifiers::{InstrumentId, Symbol, Venue},
    python::instruments::instrument_any_to_pyobject,
};
use pyo3::{prelude::*, types::PyList};
use ustr::Ustr;

use crate::{
    decode::DatabentoDecodeConfig,
    loader::DatabentoDataLoader,
    types::{DatabentoImbalance, DatabentoPublisher, DatabentoStatistics, PublisherId},
};

#[expect(clippy::needless_pass_by_value)]
#[pymethods]
#[pyo3_stub_gen::derive::gen_stub_pymethods]
impl DatabentoDataLoader {
    /// A Nautilus data loader for Databento Binary Encoding (DBN) format data.
    ///
    /// # Supported Schemas
    ///  - `MBO` -> `OrderBookDelta`
    ///  - `MBP_1` -> `(QuoteTick, Option<TradeTick>)`
    ///  - `MBP_10` -> `OrderBookDepth10`
    ///  - `BBO_1S` -> `QuoteTick`
    ///  - `BBO_1M` -> `QuoteTick`
    ///  - `CMBP_1` -> `(QuoteTick, Option<TradeTick>)`
    ///  - `CBBO_1S` -> `QuoteTick`
    ///  - `CBBO_1M` -> `QuoteTick`
    ///  - `TCBBO` -> `(QuoteTick, TradeTick)`
    ///  - `TBBO` -> `(QuoteTick, TradeTick)`
    ///  - `TRADES` -> `TradeTick`
    ///  - `OHLCV_1S` -> `Bar`
    ///  - `OHLCV_1M` -> `Bar`
    ///  - `OHLCV_1H` -> `Bar`
    ///  - `OHLCV_1D` -> `Bar`
    ///  - `OHLCV_EOD` -> `Bar`
    ///  - `DEFINITION` -> `Instrument`
    ///  - `IMBALANCE` -> `DatabentoImbalance`
    ///  - `STATISTICS` -> `DatabentoStatistics`
    ///  - `STATUS` -> `InstrumentStatus`
    ///
    /// # References
    ///
    /// <https://databento.com/docs/schemas-and-data-formats>
    #[new]
    #[pyo3(signature = (publishers_filepath=None))]
    fn py_new(publishers_filepath: Option<PathBuf>) -> PyResult<Self> {
        Self::new(publishers_filepath).map_err(to_pyvalue_err)
    }

    /// Load the publishers data from the file at the given `filepath`.
    ///
    /// # Errors
    ///
    /// Returns an error if the file cannot be read or parsed as JSON.
    #[pyo3(name = "load_publishers")]
    fn py_load_publishers(&mut self, publishers_filepath: PathBuf) -> PyResult<()> {
        self.load_publishers(publishers_filepath)
            .map_err(to_pyvalue_err)
    }

    /// Returns the internal Databento publishers currently held by the loader.
    #[must_use]
    #[pyo3(name = "get_publishers")]
    fn py_get_publishers(&self) -> HashMap<u16, DatabentoPublisher> {
        self.get_publishers()
            .iter()
            .map(|(&key, value)| (key, value.clone()))
            .collect::<HashMap<u16, DatabentoPublisher>>()
    }

    /// Sets the `venue` to map to the given `dataset`.
    #[pyo3(name = "set_dataset_for_venue")]
    fn py_set_dataset_for_venue(&mut self, dataset: String, venue: Venue) {
        self.set_dataset_for_venue(Ustr::from(&dataset), venue);
    }

    /// Returns the dataset which matches the given `venue` (if found).
    #[must_use]
    #[pyo3(name = "get_dataset_for_venue")]
    fn py_get_dataset_for_venue(&self, venue: &Venue) -> Option<String> {
        self.get_dataset_for_venue(venue).map(ToString::to_string)
    }

    /// Returns the venue which matches the given `publisher_id` (if found).
    #[must_use]
    #[pyo3(name = "get_venue_for_publisher")]
    fn py_get_venue_for_publisher(&self, publisher_id: PublisherId) -> Option<String> {
        self.get_venue_for_publisher(publisher_id)
            .map(ToString::to_string)
    }

    /// Caches a `price_precision` for the given `symbol`.
    ///
    /// When market data is read without an explicit `price_precision` argument,
    /// the loader resolves precision per record from this cache. Definitions
    /// loaded via `Self.load_instruments` are inserted automatically.
    #[pyo3(name = "set_price_precision")]
    fn py_set_price_precision(&mut self, symbol: &str, price_precision: u8) {
        self.set_price_precision(Symbol::from(symbol), price_precision);
    }

    /// Returns the cached price precisions keyed by symbol.
    #[must_use]
    #[pyo3(name = "get_price_precisions")]
    fn py_get_price_precisions(&self) -> HashMap<String, u8> {
        self.get_price_precisions()
            .iter()
            .map(|(symbol, precision)| (symbol.to_string(), *precision))
            .collect()
    }

    #[pyo3(name = "schema_for_file")]
    fn py_schema_for_file(&self, filepath: PathBuf) -> PyResult<Option<String>> {
        self.schema_from_file(&filepath).map_err(to_pyvalue_err)
    }

    /// Loads all instrument definitions from a DBN file.
    ///
    /// When `skip_on_error` is true, instruments that fail to decode are logged
    /// as warnings and skipped. When false (default), any decode error is propagated.
    ///
    /// # Errors
    ///
    /// Returns an error if loading instruments fails.
    #[pyo3(name = "load_instruments")]
    #[pyo3(signature = (filepath, use_exchange_as_venue, skip_on_error=false, expiration_overrides=None))]
    fn py_load_instruments(
        &mut self,
        py: Python,
        filepath: PathBuf,
        use_exchange_as_venue: bool,
        skip_on_error: bool,
        expiration_overrides: Option<HashMap<String, HashMap<String, String>>>,
    ) -> PyResult<Py<PyAny>> {
        let decode_config = build_decode_config(expiration_overrides)?;
        let iter = self
            .load_instruments(
                &filepath,
                use_exchange_as_venue,
                skip_on_error,
                decode_config.as_ref(),
            )
            .map_err(to_pyvalue_err)?;

        let mut data = Vec::new();

        for instrument in iter {
            let py_object = instrument_any_to_pyobject(py, instrument)?;
            data.push(py_object);
        }

        let list = PyList::new(py, &data)?;

        Ok(list.into_py_any_unwrap(py))
    }

    // Cannot include trades
    /// Loads order book delta messages from a DBN MBO schema file.
    ///
    /// Cannot include trades.
    ///
    /// # Errors
    ///
    /// Returns an error if loading order book deltas fails.
    #[pyo3(name = "load_order_book_deltas")]
    #[pyo3(signature = (filepath, instrument_id=None, price_precision=None))]
    fn py_load_order_book_deltas(
        &self,
        filepath: PathBuf,
        instrument_id: Option<InstrumentId>,
        price_precision: Option<u8>,
    ) -> PyResult<Vec<OrderBookDelta>> {
        self.load_order_book_deltas(&filepath, instrument_id, price_precision)
            .map_err(to_pyvalue_err)
    }

    /// Loads order book depth10 snapshots from a DBN MBP-10 schema file.
    ///
    /// # Errors
    ///
    /// Returns an error if loading order book depth10 fails.
    #[pyo3(name = "load_order_book_depth10")]
    #[pyo3(signature = (filepath, instrument_id=None, price_precision=None))]
    fn py_load_order_book_depth10(
        &self,
        filepath: PathBuf,
        instrument_id: Option<InstrumentId>,
        price_precision: Option<u8>,
    ) -> PyResult<Vec<OrderBookDepth10>> {
        self.load_order_book_depth10(&filepath, instrument_id, price_precision)
            .map_err(to_pyvalue_err)
    }

    /// Loads quote tick messages from a DBN MBP-1 or TBBO schema file.
    ///
    /// # Errors
    ///
    /// Returns an error if loading quotes fails.
    #[pyo3(name = "load_quotes")]
    #[pyo3(signature = (filepath, instrument_id=None, price_precision=None))]
    fn py_load_quotes(
        &self,
        filepath: PathBuf,
        instrument_id: Option<InstrumentId>,
        price_precision: Option<u8>,
    ) -> PyResult<Vec<QuoteTick>> {
        self.load_quotes(&filepath, instrument_id, price_precision)
            .map_err(to_pyvalue_err)
    }

    /// Loads best bid/offer quote messages from a DBN BBO schema file.
    ///
    /// # Errors
    ///
    /// Returns an error if loading BBO quotes fails.
    #[pyo3(name = "load_bbo_quotes")]
    #[pyo3(signature = (filepath, instrument_id=None, price_precision=None))]
    fn py_load_bbo_quotes(
        &self,
        filepath: PathBuf,
        instrument_id: Option<InstrumentId>,
        price_precision: Option<u8>,
    ) -> PyResult<Vec<QuoteTick>> {
        self.load_bbo_quotes(&filepath, instrument_id, price_precision)
            .map_err(to_pyvalue_err)
    }

    /// Loads consolidated MBP-1 quote messages from a DBN CMBP-1 schema file.
    ///
    /// # Errors
    ///
    /// Returns an error if loading consolidated MBP-1 quotes fails.
    #[pyo3(name = "load_cmbp_quotes")]
    #[pyo3(signature = (filepath, instrument_id=None, price_precision=None))]
    fn py_load_cmbp_quotes(
        &self,
        filepath: PathBuf,
        instrument_id: Option<InstrumentId>,
        price_precision: Option<u8>,
    ) -> PyResult<Vec<QuoteTick>> {
        self.load_cmbp_quotes(&filepath, instrument_id, price_precision)
            .map_err(to_pyvalue_err)
    }

    /// Loads consolidated best bid/offer quote messages from a DBN CBBO schema file.
    ///
    /// # Errors
    ///
    /// Returns an error if loading consolidated BBO quotes fails.
    #[pyo3(name = "load_cbbo_quotes")]
    #[pyo3(signature = (filepath, instrument_id=None, price_precision=None))]
    fn py_load_cbbo_quotes(
        &self,
        filepath: PathBuf,
        instrument_id: Option<InstrumentId>,
        price_precision: Option<u8>,
    ) -> PyResult<Vec<QuoteTick>> {
        self.load_cbbo_quotes(&filepath, instrument_id, price_precision)
            .map_err(to_pyvalue_err)
    }

    /// Loads trade messages from a DBN TBBO schema file.
    ///
    /// # Errors
    ///
    /// Returns an error if loading TBBO trades fails.
    #[pyo3(name = "load_tbbo_trades")]
    #[pyo3(signature = (filepath, instrument_id=None, price_precision=None))]
    fn py_load_tbbo_trades(
        &self,
        filepath: PathBuf,
        instrument_id: Option<InstrumentId>,
        price_precision: Option<u8>,
    ) -> PyResult<Vec<TradeTick>> {
        self.load_tbbo_trades(&filepath, instrument_id, price_precision)
            .map_err(to_pyvalue_err)
    }

    /// Loads trade messages from a DBN TCBBO schema file.
    ///
    /// # Errors
    ///
    /// Returns an error if loading TCBBO trades fails.
    #[pyo3(name = "load_tcbbo_trades")]
    #[pyo3(signature = (filepath, instrument_id=None, price_precision=None))]
    fn py_load_tcbbo_trades(
        &self,
        filepath: PathBuf,
        instrument_id: Option<InstrumentId>,
        price_precision: Option<u8>,
    ) -> PyResult<Vec<TradeTick>> {
        self.load_tcbbo_trades(&filepath, instrument_id, price_precision)
            .map_err(to_pyvalue_err)
    }

    /// Loads trade messages from a DBN TRADES schema file.
    ///
    /// # Errors
    ///
    /// Returns an error if loading trades fails.
    #[pyo3(name = "load_trades")]
    #[pyo3(signature = (filepath, instrument_id=None, price_precision=None))]
    fn py_load_trades(
        &self,
        filepath: PathBuf,
        instrument_id: Option<InstrumentId>,
        price_precision: Option<u8>,
    ) -> PyResult<Vec<TradeTick>> {
        self.load_trades(&filepath, instrument_id, price_precision)
            .map_err(to_pyvalue_err)
    }

    /// Loads OHLCV bar messages from a DBN OHLCV schema file.
    ///
    /// # Errors
    ///
    /// Returns an error if loading bars fails.
    #[pyo3(name = "load_bars")]
    #[pyo3(signature = (filepath, instrument_id=None, price_precision=None, timestamp_on_close=true))]
    fn py_load_bars(
        &self,
        filepath: PathBuf,
        instrument_id: Option<InstrumentId>,
        price_precision: Option<u8>,
        timestamp_on_close: bool,
    ) -> PyResult<Vec<Bar>> {
        self.load_bars(
            &filepath,
            instrument_id,
            price_precision,
            Some(timestamp_on_close),
        )
        .map_err(to_pyvalue_err)
    }

    #[pyo3(name = "load_status")]
    #[pyo3(signature = (filepath, instrument_id=None))]
    fn py_load_status(
        &self,
        filepath: PathBuf,
        instrument_id: Option<InstrumentId>,
    ) -> PyResult<Vec<InstrumentStatus>> {
        let iter = self
            .load_status_records::<dbn::StatusMsg>(&filepath, instrument_id)
            .map_err(to_pyvalue_err)?;

        let mut data = Vec::new();

        for result in iter {
            match result {
                Ok(item) => data.push(item),
                Err(e) => return Err(to_pyvalue_err(e)),
            }
        }

        Ok(data)
    }

    #[pyo3(name = "load_imbalance")]
    #[pyo3(signature = (filepath, instrument_id=None, price_precision=None))]
    fn py_load_imbalance(
        &self,
        filepath: PathBuf,
        instrument_id: Option<InstrumentId>,
        price_precision: Option<u8>,
    ) -> PyResult<Vec<DatabentoImbalance>> {
        let iter = self
            .read_imbalance_records::<dbn::ImbalanceMsg>(&filepath, instrument_id, price_precision)
            .map_err(to_pyvalue_err)?;

        let mut data = Vec::new();

        for result in iter {
            match result {
                Ok(item) => data.push(item),
                Err(e) => return Err(to_pyvalue_err(e)),
            }
        }

        Ok(data)
    }

    #[pyo3(name = "load_statistics")]
    #[pyo3(signature = (filepath, instrument_id=None, price_precision=None))]
    fn py_load_statistics(
        &self,
        filepath: PathBuf,
        instrument_id: Option<InstrumentId>,
        price_precision: Option<u8>,
    ) -> PyResult<Vec<DatabentoStatistics>> {
        let iter = self
            .read_statistics_records::<dbn::StatMsg>(&filepath, instrument_id, price_precision)
            .map_err(to_pyvalue_err)?;

        let mut data = Vec::new();

        for result in iter {
            match result {
                Ok(item) => data.push(item),
                Err(e) => return Err(to_pyvalue_err(e)),
            }
        }

        Ok(data)
    }
}

// Returns `None` when no overrides are supplied, so the loader applies its built-in defaults
fn build_decode_config(
    expiration_overrides: Option<HashMap<String, HashMap<String, String>>>,
) -> PyResult<Option<DatabentoDecodeConfig>> {
    expiration_overrides
        .map(|overrides| decode_config_from_overrides(overrides).map_err(to_pyvalue_err))
        .transpose()
}

// Builds a decode config from a dataset -> (underlying -> wall-clock time) mapping. The reserved
// underlying key "default" sets a dataset's default time; other keys are per-underlying overrides.
fn decode_config_from_overrides(
    overrides: HashMap<String, HashMap<String, String>>,
) -> Result<DatabentoDecodeConfig, String> {
    let mut config = DatabentoDecodeConfig::default();

    for (dataset_name, times) in overrides {
        let dataset = dataset_name
            .parse::<dbn::Dataset>()
            .map_err(|_| format!("Unknown dataset '{dataset_name}'"))?;
        let rule = config
            .option_expiration
            .get_mut(&dataset)
            .ok_or_else(|| format!("No expiration correction rule for dataset '{dataset_name}'"))?;

        for (underlying, value) in times {
            let time = parse_expiration_time(&value)?;
            if underlying == "default" {
                rule.default_time = time;
            } else {
                rule.overrides.insert(Ustr::from(&underlying), time);
            }
        }
    }

    Ok(config)
}

// Parses a `HH:MM` or `HH:MM:SS` wall-clock time
fn parse_expiration_time(value: &str) -> Result<Time, String> {
    value
        .parse::<Time>()
        .map_err(|_| format!("Invalid expiration time '{value}', expected 'HH:MM' or 'HH:MM:SS'"))
}

#[cfg(test)]
mod tests {
    use std::collections::HashMap;

    use databento::dbn;
    use jiff::civil::Time;
    use rstest::rstest;
    use ustr::Ustr;

    use super::decode_config_from_overrides;

    fn overrides(
        dataset: &str,
        entries: &[(&str, &str)],
    ) -> HashMap<String, HashMap<String, String>> {
        let inner = entries
            .iter()
            .map(|(k, v)| ((*k).to_string(), (*v).to_string()))
            .collect();
        HashMap::from([(dataset.to_string(), inner)])
    }

    #[rstest]
    fn test_default_key_sets_dataset_default_time() {
        let config =
            decode_config_from_overrides(overrides("OPRA.PILLAR", &[("default", "15:30")]))
                .unwrap();
        let rule = config
            .option_expiration
            .get(&dbn::Dataset::OpraPillar)
            .unwrap();
        assert_eq!(rule.default_time, Time::constant(15, 30, 0, 0));
        assert!(rule.overrides.is_empty());
    }

    #[rstest]
    fn test_underlying_key_sets_override_and_keeps_default() {
        let config =
            decode_config_from_overrides(overrides("OPRA.PILLAR", &[("SPX", "09:30:00")])).unwrap();
        let rule = config
            .option_expiration
            .get(&dbn::Dataset::OpraPillar)
            .unwrap();
        assert_eq!(
            rule.overrides.get(&Ustr::from("SPX")).copied(),
            Some(Time::constant(9, 30, 0, 0))
        );
        assert_eq!(rule.default_time, Time::constant(16, 0, 0, 0));
    }

    #[rstest]
    fn test_unknown_dataset_errors() {
        let err = decode_config_from_overrides(overrides("NOT.ADATASET", &[("default", "16:00")]))
            .unwrap_err();
        assert!(err.contains("Unknown dataset"), "was: {err}");
    }

    #[rstest]
    fn test_dataset_without_rule_errors() {
        // GLBX is a valid dataset but has no built-in expiration rule, so it cannot be tuned
        let err = decode_config_from_overrides(overrides("GLBX.MDP3", &[("default", "16:00")]))
            .unwrap_err();
        assert!(err.contains("No expiration correction rule"), "was: {err}");
    }

    #[rstest]
    fn test_invalid_time_errors() {
        let err = decode_config_from_overrides(overrides("OPRA.PILLAR", &[("default", "nope")]))
            .unwrap_err();
        assert!(err.contains("Invalid expiration time"), "was: {err}");
    }
}