1use std::str::FromStr;
19
20use nautilus_common::python::config_error_to_pyvalue_err;
21use nautilus_core::{
22 DurationNanos, UnixNanos, from_pydict,
23 python::{enums::parse_enum, params::params_to_pydict, to_pytype_err, to_pyvalue_err},
24};
25use nautilus_model::{
26 data::{NautilusDataType, NautilusRecordType},
27 instruments::NautilusInstrumentType,
28 python::{
29 common::EnumIterator,
30 data::{PyNautilusDataType, PyNautilusRecordType},
31 instruments::PyNautilusInstrumentType,
32 },
33};
34use pyo3::{
35 Bound, Py, PyAny, PyRef, PyResult, PyTypeInfo, Python,
36 types::{PyAnyMethods, PyDict, PyType},
37};
38
39use crate::config::{
40 CatalogBackendType, CatalogCompression, DEFAULT_ROTATION_TIMEZONE, DataCatalogConfig,
41 RotationConfig, RotationMode, StreamingConfig, StreamingRecordFilterConfig,
42};
43
44#[derive(Clone, Debug, Eq, Hash, PartialEq)]
45#[pyo3::pyclass(
46 frozen,
47 name = "CatalogBackend",
48 module = "nautilus_trader.persistence"
49)]
50#[pyo3_stub_gen::derive::gen_stub_pyclass(module = "nautilus_trader.persistence")]
51pub struct PyCatalogBackend {
52 inner: CatalogBackendType,
53}
54
55impl PyCatalogBackend {
56 #[must_use]
57 pub const fn new(inner: CatalogBackendType) -> Self {
58 Self { inner }
59 }
60
61 #[must_use]
62 pub fn inner(&self) -> CatalogBackendType {
63 self.inner.clone()
64 }
65}
66
67fn parse_catalog_backend(value: &str) -> PyResult<CatalogBackendType> {
68 value.parse::<CatalogBackendType>().map_err(to_pyvalue_err)
69}
70
71#[pyo3_stub_gen::derive::gen_stub_pymethods]
72#[pyo3::pymethods]
73impl PyCatalogBackend {
74 #[classattr]
75 #[expect(
76 non_snake_case,
77 clippy::use_self,
78 reason = "PyO3 stub generation needs the concrete Python enum type"
79 )]
80 fn Parquet() -> PyCatalogBackend {
81 Self::new(CatalogBackendType::Parquet)
82 }
83
84 #[staticmethod]
85 fn from_str(value: &str) -> PyResult<Self> {
86 parse_catalog_backend(value).map(Self::new)
87 }
88
89 #[staticmethod]
90 #[pyo3(name = "External")]
91 fn py_external(name: &str) -> PyResult<Self> {
92 let backend = name.parse::<CatalogBackendType>().map_err(to_pyvalue_err)?;
93 match backend {
94 CatalogBackendType::External(_) => Ok(Self::new(backend)),
95 CatalogBackendType::Parquet => Err(to_pyvalue_err(format!(
96 "Catalog backend name '{name}' is reserved for a built-in backend"
97 ))),
98 }
99 }
100
101 #[getter]
102 fn name(&self) -> &str {
103 match self.inner {
104 CatalogBackendType::Parquet => "Parquet",
105 CatalogBackendType::External(_) => "External",
106 }
107 }
108
109 #[getter]
110 fn value(&self) -> String {
111 self.inner.to_string()
112 }
113
114 #[getter]
115 fn external_name(&self) -> Option<&str> {
116 match &self.inner {
117 CatalogBackendType::External(name) => Some(name),
118 CatalogBackendType::Parquet => None,
119 }
120 }
121
122 fn __str__(&self) -> String {
123 self.inner.to_string()
124 }
125
126 fn __repr__(&self) -> String {
127 match &self.inner {
128 CatalogBackendType::Parquet => "CatalogBackend.Parquet".to_string(),
129 CatalogBackendType::External(name) => {
130 format!("CatalogBackend.External({name:?})")
131 }
132 }
133 }
134
135 fn __richcmp__(&self, other: &Self, op: pyo3::pyclass::CompareOp, py: Python<'_>) -> Py<PyAny> {
136 use nautilus_core::python::IntoPyObjectNautilusExt;
137
138 match op {
139 pyo3::pyclass::CompareOp::Eq => (self.inner == other.inner).into_py_any_unwrap(py),
140 pyo3::pyclass::CompareOp::Ne => (self.inner != other.inner).into_py_any_unwrap(py),
141 _ => py.NotImplemented(),
142 }
143 }
144
145 #[expect(
146 clippy::cast_possible_truncation,
147 clippy::cast_possible_wrap,
148 reason = "Python hashes use the platform signed integer width"
149 )]
150 fn __hash__(&self) -> isize {
151 use std::{
152 collections::hash_map::DefaultHasher,
153 hash::{Hash, Hasher},
154 };
155
156 let mut hasher = DefaultHasher::new();
157 self.inner.hash(&mut hasher);
158 hasher.finish() as isize
159 }
160}
161
162#[derive(Clone, Debug)]
163#[pyo3::pyclass(
164 frozen,
165 name = "RotationConfig",
166 module = "nautilus_trader.persistence",
167 from_py_object
168)]
169#[pyo3_stub_gen::derive::gen_stub_pyclass(module = "nautilus_trader.persistence")]
170pub struct PyRotationConfig {
171 inner: RotationConfig,
172}
173
174impl From<PyRotationConfig> for RotationConfig {
175 fn from(config: PyRotationConfig) -> Self {
176 config.inner
177 }
178}
179
180impl From<RotationConfig> for PyRotationConfig {
181 fn from(config: RotationConfig) -> Self {
182 Self { inner: config }
183 }
184}
185
186impl PyRotationConfig {
187 fn validated(inner: RotationConfig) -> PyResult<Self> {
188 inner.validate().map_err(config_error_to_pyvalue_err)?;
189 Ok(Self { inner })
190 }
191}
192
193#[pyo3_stub_gen::derive::gen_stub_pymethods]
194#[pyo3::pymethods]
195impl PyRotationConfig {
196 #[staticmethod]
197 fn no_rotation() -> Self {
198 Self {
199 inner: RotationConfig::NoRotation,
200 }
201 }
202
203 #[staticmethod]
204 fn size(max_size: u64) -> PyResult<Self> {
205 Self::validated(RotationConfig::Size { max_size })
206 }
207
208 #[staticmethod]
209 fn interval(interval_ns: u64) -> PyResult<Self> {
210 Self::validated(RotationConfig::Interval {
211 interval_ns: DurationNanos::new(interval_ns),
212 })
213 }
214
215 #[staticmethod]
216 #[pyo3(signature = (interval_ns, schedule_ns, timezone = DEFAULT_ROTATION_TIMEZONE))]
217 fn scheduled_dates(interval_ns: u64, schedule_ns: u64, timezone: &str) -> PyResult<Self> {
218 Self::validated(RotationConfig::ScheduledDates {
219 interval_ns: DurationNanos::new(interval_ns),
220 schedule_ns: UnixNanos::from(schedule_ns),
221 timezone: timezone.to_string(),
222 })
223 }
224
225 #[getter]
226 const fn mode(&self) -> RotationMode {
227 self.inner.mode()
228 }
229
230 #[getter]
231 fn max_size(&self) -> Option<u64> {
232 match self.inner {
233 RotationConfig::Size { max_size } => Some(max_size),
234 _ => None,
235 }
236 }
237
238 #[getter]
239 fn interval_ns(&self) -> Option<u64> {
240 match self.inner {
241 RotationConfig::Interval { interval_ns }
242 | RotationConfig::ScheduledDates { interval_ns, .. } => Some(interval_ns.as_u64()),
243 _ => None,
244 }
245 }
246
247 #[getter]
248 fn schedule_ns(&self) -> Option<u64> {
249 match self.inner {
250 RotationConfig::ScheduledDates { schedule_ns, .. } => Some(schedule_ns.as_u64()),
251 _ => None,
252 }
253 }
254
255 #[getter]
256 fn timezone(&self) -> Option<&str> {
257 match &self.inner {
258 RotationConfig::ScheduledDates { timezone, .. } => Some(timezone),
259 _ => None,
260 }
261 }
262
263 fn __repr__(&self) -> String {
264 format!("{:?}", self.inner)
265 }
266}
267
268#[pyo3_stub_gen::derive::gen_stub_pymethods]
269#[pyo3::pymethods]
270#[allow(
271 clippy::trivially_copy_pass_by_ref,
272 reason = "PyO3 enum methods take self by reference for Python API parity"
273)]
274impl RotationMode {
275 #[new]
277 fn py_new(py: Python<'_>, value: &Bound<'_, PyAny>) -> PyResult<Self> {
278 let t = Self::type_object(py);
279 Self::py_from_str(&t, value)
280 }
281
282 const fn __hash__(&self) -> isize {
283 *self as isize
284 }
285
286 fn __str__(&self) -> String {
287 self.to_string()
288 }
289
290 #[getter]
291 #[must_use]
292 pub fn name(&self) -> String {
293 self.to_string()
294 }
295
296 #[getter]
297 #[must_use]
298 pub const fn value(&self) -> u8 {
299 *self as u8
300 }
301
302 #[classmethod]
303 fn variants(_: &Bound<'_, PyType>, py: Python<'_>) -> EnumIterator {
304 EnumIterator::new::<Self>(py)
305 }
306
307 #[classmethod]
308 #[pyo3(name = "from_str")]
309 fn py_from_str(_: &Bound<'_, PyType>, data: &Bound<'_, PyAny>) -> PyResult<Self> {
310 let data_str: &str = data.extract()?;
311 Self::from_str(data_str).map_err(to_pyvalue_err)
312 }
313}
314
315#[pyo3_stub_gen::derive::gen_stub_pymethods]
316#[pyo3::pymethods]
317impl StreamingConfig {
318 #[new]
325 #[expect(
326 clippy::too_many_arguments,
327 clippy::fn_params_excessive_bools,
328 reason = "the PyO3 constructor mirrors the public Python configuration signature"
329 )]
330 #[pyo3(signature = (
331 writer_path,
332 catalog = None,
333 flush_interval_ms = 1000,
334 replace_existing = false,
335 rotation_config = None,
336 promotion_interval_ms = None,
337 promote_on_close = true,
338 delete_feather_after_promotion = false,
339 use_ts_event_for_ts_init = false,
340 data_types = None,
341 record_types = None,
342 instrument_types = None,
343 record_filters = None,
344 params = None,
345 ))]
346 fn py_new(
347 writer_path: String,
348 catalog: Option<DataCatalogConfig>,
349 flush_interval_ms: u64,
350 replace_existing: bool,
351 rotation_config: Option<PyRotationConfig>,
352 promotion_interval_ms: Option<u64>,
353 promote_on_close: bool,
354 delete_feather_after_promotion: bool,
355 use_ts_event_for_ts_init: bool,
356 data_types: Option<&Bound<'_, PyAny>>,
357 record_types: Option<&Bound<'_, PyAny>>,
358 instrument_types: Option<&Bound<'_, PyAny>>,
359 record_filters: Option<&Bound<'_, PyAny>>,
360 params: Option<Py<PyDict>>,
361 ) -> pyo3::PyResult<Self> {
362 let mut config = Self::new(
363 writer_path,
364 catalog,
365 flush_interval_ms,
366 replace_existing,
367 rotation_config.map_or(RotationConfig::NoRotation, Into::into),
368 );
369 config.promotion_interval_ms = promotion_interval_ms;
370 config.promote_on_close = promote_on_close;
371 config.delete_feather_after_promotion = delete_feather_after_promotion;
372 config.use_ts_event_for_ts_init = use_ts_event_for_ts_init;
373 let mut parsed_types = py_streaming_types_from_any(data_types)?;
374 parsed_types
375 .records
376 .extend(py_record_types_from_any(record_types)?.unwrap_or_default());
377 parsed_types
378 .instruments
379 .extend(py_instrument_types_from_any(instrument_types)?.unwrap_or_default());
380 config.data_types = (!parsed_types.data.is_empty()).then_some(parsed_types.data);
381 config.record_types = (!parsed_types.records.is_empty()).then_some(parsed_types.records);
382 config.instrument_types =
383 (!parsed_types.instruments.is_empty()).then_some(parsed_types.instruments);
384 config.record_filters = py_record_filters_from_any(record_filters)?;
385 config.params = Python::attach(|py| match params {
386 Some(params) => from_pydict(py, ¶ms),
387 None => Ok(None),
388 })?;
389
390 Ok(config)
391 }
392
393 #[getter]
394 fn writer_path(&self) -> &str {
395 &self.writer_path
396 }
397
398 #[getter]
399 fn catalog(&self) -> Option<DataCatalogConfig> {
400 self.catalog.clone()
401 }
402
403 #[getter]
404 const fn flush_interval_ms(&self) -> u64 {
405 self.flush_interval_ms
406 }
407
408 #[getter]
409 const fn replace_existing(&self) -> bool {
410 self.replace_existing
411 }
412
413 #[getter]
414 fn rotation_config(&self) -> PyRotationConfig {
415 self.rotation_config.clone().into()
416 }
417
418 #[getter]
419 const fn promotion_interval_ms(&self) -> Option<u64> {
420 self.promotion_interval_ms
421 }
422
423 #[getter]
424 const fn promote_on_close(&self) -> bool {
425 self.promote_on_close
426 }
427
428 #[getter]
429 const fn delete_feather_after_promotion(&self) -> bool {
430 self.delete_feather_after_promotion
431 }
432
433 #[getter]
434 const fn use_ts_event_for_ts_init(&self) -> bool {
435 self.use_ts_event_for_ts_init
436 }
437
438 #[getter]
440 #[pyo3(name = "writer_backend")]
441 fn py_writer_backend(&self) -> String {
442 self.writer_backend().to_string()
443 }
444
445 #[getter]
446 #[pyo3(name = "params")]
447 fn py_params(&self, py: Python<'_>) -> PyResult<Option<Py<PyDict>>> {
448 self.params
449 .as_ref()
450 .map(|params| params_to_pydict(py, params))
451 .transpose()
452 }
453
454 #[getter]
455 fn data_types(&self) -> Option<Vec<PyNautilusDataType>> {
456 self.data_types
457 .clone()
458 .map(|values| values.into_iter().map(PyNautilusDataType::new).collect())
459 }
460
461 #[getter]
462 fn record_types(&self) -> Option<Vec<PyNautilusRecordType>> {
463 self.record_types
464 .clone()
465 .map(|values| values.into_iter().map(PyNautilusRecordType::new).collect())
466 }
467
468 #[getter]
469 fn instrument_types(&self) -> Option<Vec<PyNautilusInstrumentType>> {
470 self.instrument_types.clone().map(|values| {
471 values
472 .into_iter()
473 .map(PyNautilusInstrumentType::new)
474 .collect()
475 })
476 }
477
478 #[getter]
479 fn record_filters(&self, py: Python<'_>) -> PyResult<Option<Py<PyDict>>> {
480 let Some(filters) = &self.record_filters else {
481 return Ok(None);
482 };
483
484 let result = PyDict::new(py);
485 for filter in filters {
486 result.set_item(
487 PyNautilusRecordType::new(filter.record_type),
488 filter.identifiers.clone(),
489 )?;
490 }
491
492 Ok(Some(result.unbind()))
493 }
494
495 fn __repr__(&self) -> String {
496 format!("{self:?}")
497 }
498}
499
500#[pyo3_stub_gen::derive::gen_stub_pymethods]
501#[pyo3::pymethods]
502impl DataCatalogConfig {
503 #[new]
505 #[expect(
506 clippy::too_many_arguments,
507 reason = "the PyO3 constructor mirrors the public Python configuration signature"
508 )]
509 #[pyo3(signature = (path, fs_protocol = None, catalog_backend = None, params = None, name = None, read_only = false, fs_rust_storage_options = None, batch_size = None, compression = None, max_row_group_size = None))]
510 fn py_new(
511 path: String,
512 fs_protocol: Option<String>,
513 catalog_backend: Option<PyRef<'_, PyCatalogBackend>>,
514 params: Option<Py<PyDict>>,
515 name: Option<String>,
516 read_only: bool,
517 fs_rust_storage_options: Option<std::collections::HashMap<String, String>>,
518 batch_size: Option<usize>,
519 compression: Option<&str>,
520 max_row_group_size: Option<usize>,
521 ) -> pyo3::PyResult<Self> {
522 let catalog_backend = catalog_backend.map(|backend| backend.inner());
523 let compression = compression
524 .map(|compression| parse_enum::<CatalogCompression>(compression, "compression"))
525 .transpose()?;
526
527 let params = Python::attach(|py| match params {
528 Some(params) => from_pydict(py, ¶ms),
529 None => Ok(None),
530 })?;
531
532 Self::builder()
533 .path(path)
534 .maybe_name(name)
535 .maybe_fs_protocol(fs_protocol)
536 .maybe_catalog_backend(catalog_backend)
537 .maybe_batch_size(batch_size)
538 .maybe_compression(compression)
539 .maybe_max_row_group_size(max_row_group_size)
540 .maybe_params(params)
541 .maybe_fs_rust_storage_options(fs_rust_storage_options.map(|options| {
542 options
543 .into_iter()
544 .map(|(key, value)| (key, value.into()))
545 .collect()
546 }))
547 .read_only(read_only)
548 .build()
549 .map_err(config_error_to_pyvalue_err)
550 }
551
552 #[getter]
554 #[pyo3(name = "path")]
555 fn py_path(&self) -> &str {
556 self.path()
557 }
558
559 #[getter]
561 #[pyo3(name = "name")]
562 fn py_name(&self) -> Option<&str> {
563 self.name()
564 }
565
566 #[getter]
568 #[pyo3(name = "read_only")]
569 fn py_read_only(&self) -> bool {
570 self.read_only()
571 }
572
573 #[getter]
575 #[pyo3(name = "fs_protocol")]
576 fn py_fs_protocol(&self) -> &str {
577 self.fs_protocol()
578 }
579
580 #[getter]
582 #[pyo3(name = "catalog_backend")]
583 fn py_catalog_backend(&self) -> PyCatalogBackend {
584 PyCatalogBackend::new(self.catalog_backend().clone())
585 }
586
587 #[getter]
589 #[pyo3(name = "batch_size")]
590 fn py_batch_size(&self) -> Option<usize> {
591 self.batch_size()
592 }
593
594 #[getter]
596 #[pyo3(name = "compression")]
597 fn py_compression(&self) -> Option<String> {
598 self.compression()
599 .map(|compression| compression.to_string())
600 }
601
602 #[getter]
604 #[pyo3(name = "max_row_group_size")]
605 fn py_max_row_group_size(&self) -> Option<usize> {
606 self.max_row_group_size()
607 }
608
609 #[getter]
611 #[pyo3(name = "params")]
612 fn py_params(&self, py: Python<'_>) -> PyResult<Option<Py<PyDict>>> {
613 self.params()
614 .map(|params| params_to_pydict(py, params))
615 .transpose()
616 }
617
618 #[getter]
619 fn fs_rust_storage_option_keys(&self) -> Option<Vec<String>> {
620 self.fs_rust_storage_options().map(|options| {
621 let mut keys = options.keys().cloned().collect::<Vec<_>>();
622 keys.sort();
623 keys
624 })
625 }
626
627 fn __repr__(&self) -> String {
628 format!("{self:?}")
629 }
630}
631
632#[derive(Default)]
633struct ParsedStreamingTypes {
634 data: Vec<NautilusDataType>,
635 records: Vec<NautilusRecordType>,
636 instruments: Vec<NautilusInstrumentType>,
637}
638
639fn py_streaming_type_from_any(
640 value: &Bound<'_, PyAny>,
641 parsed: &mut ParsedStreamingTypes,
642) -> pyo3::PyResult<()> {
643 if let Ok(data_type) = value.extract::<PyRef<'_, PyNautilusDataType>>() {
644 parsed.data.push(data_type.inner());
645 return Ok(());
646 }
647
648 if let Ok(record_type) = value.extract::<PyRef<'_, PyNautilusRecordType>>() {
649 parsed.records.push(record_type.inner());
650 return Ok(());
651 }
652
653 if let Ok(instrument_type) = value.extract::<PyRef<'_, PyNautilusInstrumentType>>() {
654 parsed.instruments.push(instrument_type.inner());
655 return Ok(());
656 }
657
658 Err(to_pytype_err(
659 "streaming type must be NautilusDataType, NautilusRecordType, or NautilusInstrumentType",
660 ))
661}
662
663fn py_streaming_types_from_any(
664 values: Option<&Bound<'_, PyAny>>,
665) -> pyo3::PyResult<ParsedStreamingTypes> {
666 let mut parsed = ParsedStreamingTypes::default();
667
668 if let Some(values) = values {
669 for item in values.try_iter()? {
670 py_streaming_type_from_any(&item?, &mut parsed)?;
671 }
672 }
673
674 Ok(parsed)
675}
676
677fn py_record_type_from_any(record_type: &Bound<'_, PyAny>) -> pyo3::PyResult<NautilusRecordType> {
678 if let Ok(record_type) = record_type.extract::<PyRef<'_, PyNautilusRecordType>>() {
679 return Ok(record_type.inner());
680 }
681
682 Err(to_pytype_err("record_type must be NautilusRecordType"))
683}
684
685fn py_record_types_from_any(
686 record_types: Option<&Bound<'_, PyAny>>,
687) -> pyo3::PyResult<Option<Vec<NautilusRecordType>>> {
688 record_types
689 .map(|record_types| {
690 record_types
691 .try_iter()?
692 .map(|item| py_record_type_from_any(&item?))
693 .collect::<pyo3::PyResult<Vec<_>>>()
694 })
695 .transpose()
696}
697
698pub(crate) fn py_instrument_type_from_any(
699 instrument_type: &Bound<'_, PyAny>,
700) -> pyo3::PyResult<NautilusInstrumentType> {
701 if let Ok(instrument_type) = instrument_type.extract::<PyRef<'_, PyNautilusInstrumentType>>() {
702 return Ok(instrument_type.inner());
703 }
704
705 Err(to_pytype_err(
706 "instrument_type must be NautilusInstrumentType",
707 ))
708}
709
710fn py_instrument_types_from_any(
711 instrument_types: Option<&Bound<'_, PyAny>>,
712) -> pyo3::PyResult<Option<Vec<NautilusInstrumentType>>> {
713 instrument_types
714 .map(|instrument_types| {
715 instrument_types
716 .try_iter()?
717 .map(|item| py_instrument_type_from_any(&item?))
718 .collect::<pyo3::PyResult<Vec<_>>>()
719 })
720 .transpose()
721}
722
723fn py_record_filters_from_any(
724 record_filters: Option<&Bound<'_, PyAny>>,
725) -> pyo3::PyResult<Option<Vec<StreamingRecordFilterConfig>>> {
726 let Some(record_filters) = record_filters else {
727 return Ok(None);
728 };
729
730 let record_filters = record_filters.cast::<PyDict>()?;
731 let mut filters = Vec::with_capacity(record_filters.len()?);
732 for (record_type, identifiers) in record_filters {
733 let record_type = py_record_type_from_any(&record_type)?;
734
735 let identifiers = if identifiers.is_none() {
736 None
737 } else if let Ok(identifier) = identifiers.extract::<String>() {
738 Some(vec![identifier])
739 } else {
740 Some(identifiers.extract::<Vec<String>>()?)
741 };
742
743 filters.push(StreamingRecordFilterConfig {
744 record_type,
745 identifiers,
746 });
747 }
748
749 Ok((!filters.is_empty()).then_some(filters))
750}
751
752#[cfg(test)]
753mod tests {
754 use ahash::AHashMap;
755 use rstest::rstest;
756
757 use super::*;
758
759 #[rstest]
760 fn test_data_catalog_config_py_new_keeps_storage_option_values() {
761 Python::initialize();
762
763 let config = DataCatalogConfig::py_new(
764 "bucket/catalog".to_string(),
765 Some("s3".to_string()),
766 None,
767 None,
768 None,
769 false,
770 Some(std::collections::HashMap::from([(
771 "aws_secret_access_key".to_string(),
772 "catalog-secret".to_string(),
773 )])),
774 None,
775 None,
776 None,
777 )
778 .unwrap();
779
780 assert_eq!(
781 config.connect_config().storage_options,
782 Some(AHashMap::from([(
783 "aws_secret_access_key".to_string(),
784 "catalog-secret".to_string(),
785 )]))
786 );
787 }
788}