nautilus_model/data/
registry.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//! Registries for custom data: JSON (de)serialization and Arrow encode/decode.
17//!
18//! Mirrors Python's `register_serializable_type` and `register_arrow` in `custom.py`.
19//! The registry only stores type name -> callbacks for lookup; each type provides
20//! its own deserialize/encode/decode via the trait or registration.
21
22use std::sync::Arc;
23
24#[cfg(feature = "arrow")]
25use arrow::{datatypes::Schema, record_batch::RecordBatch};
26use dashmap::{DashMap, mapref::entry::Entry};
27use nautilus_core::Params;
28#[cfg(feature = "python")]
29use pyo3::types::PyAnyMethods;
30
31use crate::data::{CustomData, CustomDataTrait, Data, DataType};
32
33pub type JsonDeserializer =
34    Box<dyn Fn(serde_json::Value) -> Result<Arc<dyn CustomDataTrait>, anyhow::Error> + Send + Sync>;
35#[cfg(feature = "arrow")]
36pub type ArrowEncoder =
37    Box<dyn Fn(&[Arc<dyn CustomDataTrait>]) -> Result<RecordBatch, anyhow::Error> + Send + Sync>;
38#[cfg(feature = "arrow")]
39pub type ArrowDecoder = Box<
40    dyn Fn(
41            &std::collections::HashMap<String, String>,
42            RecordBatch,
43        ) -> Result<Vec<Data>, anyhow::Error>
44        + Send
45        + Sync,
46>;
47
48struct Registries {
49    json: DashMap<String, JsonDeserializer>,
50    #[cfg(feature = "arrow")]
51    arrow: DashMap<String, (Arc<Schema>, ArrowEncoder, ArrowDecoder)>,
52}
53
54fn registries() -> &'static Registries {
55    static REGISTRIES: std::sync::OnceLock<Registries> = std::sync::OnceLock::new();
56    REGISTRIES.get_or_init(|| Registries {
57        json: DashMap::new(),
58        #[cfg(feature = "arrow")]
59        arrow: DashMap::new(),
60    })
61}
62
63/// Registers a JSON deserializer for the given custom data type name.
64/// When `Data::deserialize` sees this type name, it will call this function.
65///
66/// # Errors
67/// Returns an error if the type is already registered.
68pub fn register_json_deserializer(
69    type_name: &str,
70    deserializer: JsonDeserializer,
71) -> Result<(), anyhow::Error> {
72    let reg = registries();
73    match reg.json.entry(type_name.to_string()) {
74        Entry::Occupied(_) => {
75            anyhow::bail!("Custom data type \"{type_name}\" is already registered for JSON");
76        }
77        Entry::Vacant(v) => {
78            v.insert(deserializer);
79            Ok(())
80        }
81    }
82}
83
84/// Registers a JSON deserializer for the given custom data type name if not already registered.
85/// If the type is already registered, returns `Ok(())` without overwriting (idempotent).
86/// Use this where repeated registration can occur (e.g. module init).
87///
88/// # Errors
89/// Does not return an error (idempotent insert into `DashMap`).
90pub fn ensure_json_deserializer_registered(
91    type_name: &str,
92    deserializer: JsonDeserializer,
93) -> Result<(), anyhow::Error> {
94    let reg = registries();
95    reg.json
96        .entry(type_name.to_string())
97        .or_insert_with(|| deserializer);
98    Ok(())
99}
100
101/// Parses a "`data_type`" JSON object into `DataType` (`type_name`, metadata, identifier).
102fn parse_data_type_from_value(value: &serde_json::Value) -> Option<DataType> {
103    let obj = value.get("data_type")?.as_object()?;
104    let type_name = obj.get("type_name")?.as_str()?;
105    let metadata = obj.get("metadata").and_then(|m| {
106        if m.is_null() {
107            None
108        } else {
109            let p: Params = serde_json::from_value(m.clone()).ok()?;
110            if p.is_empty() { None } else { Some(p) }
111        }
112    });
113    let identifier = obj
114        .get("identifier")
115        .and_then(|v| v.as_str())
116        .map(String::from);
117    Some(DataType::new(type_name, metadata, identifier))
118}
119
120/// Parses the canonical `CustomData` JSON envelope `{ type, data_type, payload }` and returns
121/// the payload value to pass to the registered type deserializer. Does not depend on
122/// user payload field names.
123fn parse_envelope_payload(value: &serde_json::Value) -> Result<serde_json::Value, anyhow::Error> {
124    let payload = value
125        .get("payload")
126        .cloned()
127        .ok_or_else(|| anyhow::anyhow!("CustomData JSON missing 'payload' field"))?;
128    Ok(payload)
129}
130
131/// Looks up and runs the JSON deserializer for the given type name.
132/// Returns `None` if the type is not registered.
133///
134/// # Errors
135/// Returns an error if the deserializer fails.
136pub fn deserialize_custom_from_json(
137    type_name: &str,
138    value: &serde_json::Value,
139) -> Result<Option<Data>, anyhow::Error> {
140    let reg = registries();
141    let deserializer_ref = match reg.json.get(type_name) {
142        Some(d) => d,
143        None => return Ok(None),
144    };
145    let data_type = parse_data_type_from_value(value);
146    let payload = parse_envelope_payload(value)?;
147    let arc = deserializer_ref.value()(payload)?;
148    let custom = match data_type {
149        Some(dt) => CustomData::new(arc, dt),
150        None => CustomData::from_arc(arc),
151    };
152    Ok(Some(Data::Custom(custom)))
153}
154
155/// Registers Arrow schema, encoder, and decoder for the given custom data type name.
156///
157/// # Errors
158/// Returns an error if the type is already registered for Arrow.
159#[cfg(feature = "arrow")]
160pub fn register_arrow(
161    type_name: &str,
162    schema: Arc<Schema>,
163    encoder: ArrowEncoder,
164    decoder: ArrowDecoder,
165) -> Result<(), anyhow::Error> {
166    let reg = registries();
167    match reg.arrow.entry(type_name.to_string()) {
168        Entry::Occupied(_) => {
169            anyhow::bail!("Custom data type \"{type_name}\" is already registered for Arrow");
170        }
171        Entry::Vacant(v) => {
172            v.insert((schema, encoder, decoder));
173            Ok(())
174        }
175    }
176}
177
178/// Registers Arrow schema, encoder, and decoder for the given custom data type name if not already
179/// registered. If the type is already registered, returns `Ok(())` without overwriting (idempotent).
180/// Use this where repeated registration can occur (e.g. module init).
181///
182/// # Errors
183/// Does not return an error (idempotent insert into `DashMap`).
184#[cfg(feature = "arrow")]
185pub fn ensure_arrow_registered(
186    type_name: &str,
187    schema: Arc<Schema>,
188    encoder: ArrowEncoder,
189    decoder: ArrowDecoder,
190) -> Result<(), anyhow::Error> {
191    let reg = registries();
192    reg.arrow
193        .entry(type_name.to_string())
194        .or_insert_with(|| (schema, encoder, decoder));
195    Ok(())
196}
197
198/// Returns the Arrow schema for the given custom type name, if registered.
199#[must_use]
200#[cfg(feature = "arrow")]
201pub fn get_arrow_schema(type_name: &str) -> Option<Arc<Schema>> {
202    let reg = registries();
203    reg.arrow
204        .get(type_name)
205        .map(|entry| Arc::clone(&entry.value().0))
206}
207
208/// Encodes a slice of custom data trait objects to a `RecordBatch` using the registered encoder.
209///
210/// # Errors
211/// Returns an error if the type is not registered or encoding fails.
212#[cfg(feature = "arrow")]
213pub fn encode_custom_to_arrow(
214    type_name: &str,
215    items: &[Arc<dyn CustomDataTrait>],
216) -> Result<Option<RecordBatch>, anyhow::Error> {
217    let reg = registries();
218    let entry = match reg.arrow.get(type_name) {
219        Some(e) => e,
220        None => return Ok(None),
221    };
222    let encoder = &entry.value().1;
223    encoder(items).map(Some)
224}
225
226/// Decodes a `RecordBatch` into `Vec<Data>` using the registered decoder.
227///
228/// # Errors
229/// Returns an error if the type is not registered or decoding fails.
230#[expect(
231    clippy::implicit_hasher,
232    reason = "callers always use the default hasher"
233)]
234#[cfg(feature = "arrow")]
235pub fn decode_custom_from_arrow(
236    type_name: &str,
237    metadata: &std::collections::HashMap<String, String>,
238    record_batch: RecordBatch,
239) -> Result<Option<Vec<Data>>, anyhow::Error> {
240    let reg = registries();
241    let entry = match reg.arrow.get(type_name) {
242        Some(e) => e,
243        None => return Ok(None),
244    };
245    let decoder = &entry.value().2;
246    decoder(metadata, record_batch).map(Some)
247}
248
249#[cfg(feature = "python")]
250pub type PyExtractor = Box<
251    dyn for<'a> Fn(&pyo3::Bound<'a, pyo3::PyAny>) -> Option<Arc<dyn CustomDataTrait>> + Send + Sync,
252>;
253
254#[cfg(feature = "python")]
255fn py_extractors() -> &'static DashMap<String, PyExtractor> {
256    static PY_EXTRACTORS: std::sync::OnceLock<DashMap<String, PyExtractor>> =
257        std::sync::OnceLock::new();
258    PY_EXTRACTORS.get_or_init(DashMap::new)
259}
260
261/// Registers a `PyExtractor` for the given custom data type name.
262/// Used by `CustomData` constructor to convert Python objects to `Arc<dyn CustomDataTrait>`.
263///
264/// # Errors
265/// Returns an error if the type is already registered.
266#[cfg(feature = "python")]
267pub fn register_py_extractor(type_name: &str, extractor: PyExtractor) -> Result<(), anyhow::Error> {
268    let reg = py_extractors();
269    match reg.entry(type_name.to_string()) {
270        Entry::Occupied(_) => {
271            anyhow::bail!(
272                "Custom data type \"{type_name}\" is already registered for Python extraction"
273            );
274        }
275        Entry::Vacant(v) => {
276            v.insert(extractor);
277            Ok(())
278        }
279    }
280}
281
282/// Registers a `PyExtractor` for the given custom data type name if not already registered.
283/// If the type is already registered, returns `Ok(())` without overwriting (idempotent).
284/// Use this where repeated registration can occur (e.g. module init).
285///
286/// # Errors
287/// Does not return an error (idempotent insert into `DashMap`).
288#[cfg(feature = "python")]
289pub fn ensure_py_extractor_registered(
290    type_name: &str,
291    extractor: PyExtractor,
292) -> Result<(), anyhow::Error> {
293    let reg = py_extractors();
294    reg.entry(type_name.to_string())
295        .or_insert_with(|| extractor);
296    Ok(())
297}
298
299/// Tries to extract `Arc<dyn CustomDataTrait>` from a Python object using the registered extractor.
300/// Returns None if no extractor is registered or extraction fails.
301#[cfg(feature = "python")]
302#[must_use]
303pub fn try_extract_from_py(
304    type_name: &str,
305    obj: &pyo3::Bound<'_, pyo3::PyAny>,
306) -> Option<Arc<dyn CustomDataTrait>> {
307    let reg = py_extractors();
308    let entry = reg.get(type_name)?;
309    let extractor = entry.value();
310    extractor(obj)
311}
312
313#[cfg(feature = "python")]
314type RustExtractorFactory = Box<dyn Fn() -> PyExtractor + Send + Sync>;
315
316#[cfg(feature = "python")]
317fn rust_extractor_factories() -> &'static DashMap<String, RustExtractorFactory> {
318    static RUST_EXTRACTOR_FACTORIES: std::sync::OnceLock<DashMap<String, RustExtractorFactory>> =
319        std::sync::OnceLock::new();
320    RUST_EXTRACTOR_FACTORIES.get_or_init(DashMap::new)
321}
322
323/// Registers a factory that produces a `PyExtractor` for the given type name.
324/// Crates (e.g. persistence) call this at load time for each Rust custom data type.
325/// When `register_custom_data_class(cls)` is called with that type's class, the factory is invoked
326/// and the extractor is registered in the main `PyExtractor` registry.
327///
328/// # Errors
329/// Returns an error if the type name is already registered.
330#[cfg(feature = "python")]
331pub fn register_rust_extractor_factory(
332    type_name: &str,
333    factory: RustExtractorFactory,
334) -> Result<(), anyhow::Error> {
335    let reg = rust_extractor_factories();
336    match reg.entry(type_name.to_string()) {
337        Entry::Occupied(_) => {
338            anyhow::bail!("Rust extractor factory for \"{type_name}\" is already registered");
339        }
340        Entry::Vacant(v) => {
341            v.insert(factory);
342            Ok(())
343        }
344    }
345}
346
347/// Registers a factory that produces a `PyExtractor` for the given type name if not already
348/// registered. If the type is already registered, returns `Ok(())` without overwriting (idempotent).
349/// Use this where repeated registration can occur (e.g. module load).
350///
351/// # Errors
352/// Does not return an error (idempotent insert into `DashMap`).
353#[cfg(feature = "python")]
354pub fn ensure_rust_extractor_factory_registered(
355    type_name: &str,
356    factory: RustExtractorFactory,
357) -> Result<(), anyhow::Error> {
358    let reg = rust_extractor_factories();
359    reg.entry(type_name.to_string()).or_insert_with(|| factory);
360    Ok(())
361}
362
363/// Registers a Rust custom data type for Python extraction. Call once per type at module load
364/// (e.g. in the persistence PyO3 module). Uses [`register_rust_extractor_factory`] with a
365/// factory that builds the extractor for `T`.
366///
367/// # Errors
368/// Returns an error if the type name is already registered.
369#[cfg(feature = "python")]
370pub fn register_rust_extractor<T>() -> Result<(), anyhow::Error>
371where
372    T: CustomDataTrait + for<'a, 'py> pyo3::FromPyObject<'a, 'py> + Send + Sync + 'static,
373{
374    let type_name = T::type_name_static();
375    let factory: RustExtractorFactory = Box::new(|| {
376        Box::new(|obj: &pyo3::Bound<'_, pyo3::PyAny>| {
377            obj.extract::<T>()
378                .ok()
379                .map(|x| Arc::new(x) as Arc<dyn CustomDataTrait>)
380        })
381    });
382    register_rust_extractor_factory(type_name, factory)
383}
384
385/// Registers a Rust custom data type for Python extraction if not already registered.
386/// If the type is already registered, returns `Ok(())` without overwriting (idempotent).
387/// Use this where repeated registration can occur (e.g. module load).
388///
389/// # Errors
390/// Does not return an error (idempotent insert into `DashMap`).
391#[cfg(feature = "python")]
392pub fn ensure_rust_extractor_registered<T>() -> Result<(), anyhow::Error>
393where
394    T: CustomDataTrait + for<'a, 'py> pyo3::FromPyObject<'a, 'py> + Send + Sync + 'static,
395{
396    let type_name = T::type_name_static();
397    let factory: RustExtractorFactory = Box::new(|| {
398        Box::new(|obj: &pyo3::Bound<'_, pyo3::PyAny>| {
399            obj.extract::<T>()
400                .ok()
401                .map(|x| Arc::new(x) as Arc<dyn CustomDataTrait>)
402        })
403    });
404    ensure_rust_extractor_factory_registered(type_name, factory)
405}
406
407/// Calls the registered factory for the given type name and returns the extractor, if any.
408#[cfg(feature = "python")]
409#[must_use]
410pub fn get_rust_extractor(type_name: &str) -> Option<PyExtractor> {
411    let reg = rust_extractor_factories();
412    let factory_ref = reg.get(type_name)?;
413    Some(factory_ref.value()())
414}
415
416#[cfg(test)]
417mod tests {
418    use nautilus_core::UnixNanos;
419    use rstest::rstest;
420    use serde::{Deserialize, Serialize};
421
422    use super::*;
423    use crate::data::{CustomData, custom::register_custom_data_json};
424
425    #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
426    struct TestRegCustomData {
427        ts_init: UnixNanos,
428    }
429
430    #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
431    #[serde(deny_unknown_fields)]
432    struct StrictRegCustomData {
433        ts_init: UnixNanos,
434    }
435
436    impl crate::data::HasTsInit for TestRegCustomData {
437        fn ts_init(&self) -> UnixNanos {
438            self.ts_init
439        }
440    }
441
442    impl crate::data::custom::CustomDataTrait for TestRegCustomData {
443        fn type_name(&self) -> &'static str {
444            "TestRegCustomData"
445        }
446        fn type_name_static() -> &'static str {
447            "TestRegCustomData"
448        }
449        fn as_any(&self) -> &dyn std::any::Any {
450            self
451        }
452        fn ts_event(&self) -> nautilus_core::UnixNanos {
453            self.ts_init
454        }
455        fn to_json(&self) -> anyhow::Result<String> {
456            Ok(serde_json::to_string(self)?)
457        }
458        fn clone_arc(&self) -> Arc<dyn crate::data::CustomDataTrait> {
459            Arc::new(self.clone())
460        }
461        fn eq_arc(&self, other: &dyn crate::data::CustomDataTrait) -> bool {
462            other.as_any().downcast_ref::<Self>() == Some(self)
463        }
464        fn from_json(
465            value: serde_json::Value,
466        ) -> anyhow::Result<Arc<dyn crate::data::CustomDataTrait>> {
467            let t: Self = serde_json::from_value(value)?;
468            Ok(Arc::new(t))
469        }
470    }
471
472    impl crate::data::HasTsInit for StrictRegCustomData {
473        fn ts_init(&self) -> UnixNanos {
474            self.ts_init
475        }
476    }
477
478    impl crate::data::custom::CustomDataTrait for StrictRegCustomData {
479        fn type_name(&self) -> &'static str {
480            "StrictRegCustomData"
481        }
482        fn type_name_static() -> &'static str {
483            "StrictRegCustomData"
484        }
485        fn as_any(&self) -> &dyn std::any::Any {
486            self
487        }
488        fn ts_event(&self) -> nautilus_core::UnixNanos {
489            self.ts_init
490        }
491        fn to_json(&self) -> anyhow::Result<String> {
492            Ok(serde_json::to_string(self)?)
493        }
494        fn clone_arc(&self) -> Arc<dyn crate::data::CustomDataTrait> {
495            Arc::new(self.clone())
496        }
497        fn eq_arc(&self, other: &dyn crate::data::CustomDataTrait) -> bool {
498            other.as_any().downcast_ref::<Self>() == Some(self)
499        }
500        fn from_json(
501            value: serde_json::Value,
502        ) -> anyhow::Result<Arc<dyn crate::data::CustomDataTrait>> {
503            let t: Self = serde_json::from_value(value)?;
504            Ok(Arc::new(t))
505        }
506    }
507
508    #[rstest]
509    fn json_registry_roundtrip() {
510        let _ = register_custom_data_json::<TestRegCustomData>();
511
512        let data = Data::Custom(CustomData::from_arc(Arc::new(TestRegCustomData {
513            ts_init: UnixNanos::from(100),
514        })));
515
516        let json = serde_json::to_string(&data).unwrap();
517        let back: Data = serde_json::from_str(&json).unwrap();
518
519        match (&data, &back) {
520            (Data::Custom(a), Data::Custom(b)) => {
521                assert_eq!(a.data.type_name(), b.data.type_name());
522                assert_eq!(a.data.ts_init(), b.data.ts_init());
523            }
524            _ => panic!("expected Custom variant"),
525        }
526    }
527
528    #[rstest]
529    fn json_registry_roundtrip_with_deny_unknown_fields() {
530        let _ = register_custom_data_json::<StrictRegCustomData>();
531
532        let data = Data::Custom(CustomData::from_arc(Arc::new(StrictRegCustomData {
533            ts_init: UnixNanos::from(200),
534        })));
535
536        let json = serde_json::to_string(&data).unwrap();
537        let back: Data = serde_json::from_str(&json).unwrap();
538
539        match (&data, &back) {
540            (Data::Custom(a), Data::Custom(b)) => {
541                assert_eq!(a.data.type_name(), b.data.type_name());
542                assert_eq!(a.data.ts_init(), b.data.ts_init());
543            }
544            _ => panic!("expected Custom variant"),
545        }
546    }
547
548    #[rstest]
549    fn ensure_json_deserializer_registered_is_idempotent() {
550        let deserializer: JsonDeserializer = Box::new(|value| {
551            let t: TestRegCustomData = serde_json::from_value(value)?;
552            Ok(Arc::new(t) as Arc<dyn crate::data::CustomDataTrait>)
553        });
554        let r1 = ensure_json_deserializer_registered("IdempotentTestJson", deserializer);
555        assert!(r1.is_ok(), "first registration should succeed");
556        let deserializer2: JsonDeserializer = Box::new(|value| {
557            let t: TestRegCustomData = serde_json::from_value(value)?;
558            Ok(Arc::new(t) as Arc<dyn crate::data::CustomDataTrait>)
559        });
560        let r2 = ensure_json_deserializer_registered("IdempotentTestJson", deserializer2);
561        assert!(
562            r2.is_ok(),
563            "second registration with same type_name should succeed (idempotent)"
564        );
565    }
566
567    #[rstest]
568    fn register_json_deserializer_fails_on_duplicate() {
569        let deserializer: JsonDeserializer = Box::new(|value| {
570            let t: TestRegCustomData = serde_json::from_value(value)?;
571            Ok(Arc::new(t) as Arc<dyn crate::data::CustomDataTrait>)
572        });
573        let r1 = register_json_deserializer("StrictDuplicateTestJson", deserializer);
574        assert!(r1.is_ok());
575        let deserializer2: JsonDeserializer = Box::new(|value| {
576            let t: TestRegCustomData = serde_json::from_value(value)?;
577            Ok(Arc::new(t) as Arc<dyn crate::data::CustomDataTrait>)
578        });
579        let r2 = register_json_deserializer("StrictDuplicateTestJson", deserializer2);
580        assert!(r2.is_err());
581        let err_msg = r2.unwrap_err().to_string();
582        assert!(
583            err_msg.contains("already registered"),
584            "expected 'already registered' in error, found: {err_msg}"
585        );
586    }
587
588    #[rstest]
589    #[cfg(feature = "arrow")]
590    fn ensure_arrow_registered_is_idempotent() {
591        let schema = Arc::new(arrow::datatypes::Schema::empty());
592        let encoder: ArrowEncoder = Box::new(|_| {
593            Ok(arrow::record_batch::RecordBatch::new_empty(Arc::new(
594                arrow::datatypes::Schema::empty(),
595            )))
596        });
597        let decoder: ArrowDecoder = Box::new(|_, _| Ok(Vec::new()));
598
599        let r1 = ensure_arrow_registered("IdempotentTestArrow", schema, encoder, decoder);
600        assert!(r1.is_ok(), "first Arrow registration should succeed");
601
602        let schema2 = Arc::new(arrow::datatypes::Schema::empty());
603        let encoder2: ArrowEncoder = Box::new(|_| {
604            Ok(arrow::record_batch::RecordBatch::new_empty(Arc::new(
605                arrow::datatypes::Schema::empty(),
606            )))
607        });
608        let decoder2: ArrowDecoder = Box::new(|_, _| Ok(Vec::new()));
609
610        let r2 = ensure_arrow_registered("IdempotentTestArrow", schema2, encoder2, decoder2);
611        assert!(
612            r2.is_ok(),
613            "second Arrow registration with same type_name should be idempotent"
614        );
615    }
616}
nautilus_model/data/registry.rs

nautilus_model/data/
registry.rs
