1#[cfg(feature = "python")]
17use std::collections::HashSet;
18use std::{any::Any, fmt::Debug, sync::Arc};
19
20use nautilus_core::UnixNanos;
21#[cfg(feature = "python")]
22use parking_lot::RwLock;
23#[cfg(feature = "python")]
24use pyo3::{IntoPyObjectExt, prelude::*, types::PyAny};
25use serde::{Serialize, Serializer};
26
27use crate::data::{
28 Data, DataType, HasTsInit,
29 registry::{ensure_json_deserializer_registered, register_json_deserializer},
30};
31
32#[cfg(feature = "python")]
33fn intern_type_name_static(name: String) -> &'static str {
34 static INTERNER: std::sync::OnceLock<RwLock<HashSet<&'static str>>> =
35 std::sync::OnceLock::new();
36 let set = INTERNER.get_or_init(|| RwLock::new(HashSet::new()));
37
38 let guard = set.read();
39 if let Some(&existing) = guard.get(name.as_str()) {
40 return existing;
41 }
42 drop(guard);
43
44 let mut guard = set.write();
45 if let Some(&existing) = guard.get(name.as_str()) {
46 return existing;
47 }
48 let leaked: &'static str = Box::leak(name.into_boxed_str());
49 guard.insert(leaked);
50 leaked
51}
52
53#[cfg(feature = "python")]
60pub struct PythonCustomDataWrapper {
61 py_object: Py<PyAny>,
63 cached_ts_event: UnixNanos,
65 cached_ts_init: UnixNanos,
67 cached_type_name: String,
69 cached_type_name_static: &'static str,
71}
72
73#[cfg(feature = "python")]
74impl PythonCustomDataWrapper {
75 pub fn new(_py: Python<'_>, py_object: &Bound<'_, PyAny>) -> PyResult<Self> {
82 let ts_event: u64 = py_object.getattr("ts_event")?.extract()?;
84 let ts_event = UnixNanos::from(ts_event);
85
86 let ts_init: u64 = py_object.getattr("ts_init")?.extract()?;
88 let ts_init = UnixNanos::from(ts_init);
89
90 let data_class = py_object.get_type();
92 let type_name: String = if data_class.hasattr("type_name_static")? {
93 data_class.call_method0("type_name_static")?.extract()?
94 } else {
95 data_class.getattr("__name__")?.extract()?
96 };
97
98 let type_name_static: &'static str = intern_type_name_static(type_name.clone());
100
101 Ok(Self {
102 py_object: py_object.clone().unbind(),
103 cached_ts_event: ts_event,
104 cached_ts_init: ts_init,
105 cached_type_name: type_name,
106 cached_type_name_static: type_name_static,
107 })
108 }
109
110 #[must_use]
112 pub fn py_object(&self) -> &Py<PyAny> {
113 &self.py_object
114 }
115
116 #[must_use]
118 pub fn get_type_name(&self) -> &str {
119 &self.cached_type_name
120 }
121}
122
123#[cfg(feature = "python")]
124impl Clone for PythonCustomDataWrapper {
125 fn clone(&self) -> Self {
126 Python::attach(|py| Self {
127 py_object: self.py_object.clone_ref(py),
128 cached_ts_event: self.cached_ts_event,
129 cached_ts_init: self.cached_ts_init,
130 cached_type_name: self.cached_type_name.clone(),
131 cached_type_name_static: self.cached_type_name_static,
132 })
133 }
134}
135
136#[cfg(feature = "python")]
137impl Debug for PythonCustomDataWrapper {
138 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
139 f.debug_struct(stringify!(PythonCustomDataWrapper))
140 .field("py_object", &self.py_object)
141 .field("type_name", &self.cached_type_name)
142 .field("type_name_static", &self.cached_type_name_static)
143 .field("ts_event", &self.cached_ts_event)
144 .field("ts_init", &self.cached_ts_init)
145 .finish()
146 }
147}
148
149#[cfg(feature = "python")]
150impl HasTsInit for PythonCustomDataWrapper {
151 fn ts_init(&self) -> UnixNanos {
152 self.cached_ts_init
153 }
154}
155
156#[cfg(feature = "python")]
157impl CustomDataTrait for PythonCustomDataWrapper {
158 fn type_name(&self) -> &'static str {
159 self.cached_type_name_static
160 }
161
162 fn as_any(&self) -> &dyn Any {
163 self
164 }
165
166 fn ts_event(&self) -> UnixNanos {
167 self.cached_ts_event
168 }
169
170 fn to_json(&self) -> anyhow::Result<String> {
171 Python::attach(|py| {
172 let obj = self.py_object.bind(py);
173 if obj.hasattr("to_json")? {
175 let json_str: String = obj.call_method0("to_json")?.extract()?;
176 Ok(json_str)
177 } else {
178 let json_module = py.import("json")?;
180 let dict = if obj.hasattr("__dict__")? {
182 obj.getattr("__dict__")?
183 } else {
184 anyhow::bail!("Python object has no to_json() method or __dict__ attribute");
185 };
186 let json_str: String = json_module.call_method1("dumps", (dict,))?.extract()?;
187 Ok(json_str)
188 }
189 })
190 }
191
192 fn clone_arc(&self) -> Arc<dyn CustomDataTrait> {
193 Arc::new(self.clone())
194 }
195
196 fn eq_arc(&self, other: &dyn CustomDataTrait) -> bool {
197 if let Some(other_wrapper) = other.as_any().downcast_ref::<Self>() {
200 Python::attach(|py| {
201 let a = self.py_object.bind(py);
202 let b = other_wrapper.py_object.bind(py);
203 if a.is(b) {
204 return true;
205 }
206 a.eq(b).unwrap_or(false)
207 })
208 } else {
209 false
210 }
211 }
212
213 fn to_pyobject(&self, py: Python<'_>) -> PyResult<Py<PyAny>> {
214 Ok(self.py_object.clone_ref(py))
216 }
217}
218
219#[cfg(feature = "python")]
220fn python_data_classes() -> &'static dashmap::DashMap<String, Py<PyAny>> {
221 static PYTHON_DATA_CLASSES: std::sync::OnceLock<dashmap::DashMap<String, Py<PyAny>>> =
222 std::sync::OnceLock::new();
223 PYTHON_DATA_CLASSES.get_or_init(dashmap::DashMap::new)
224}
225
226#[cfg(feature = "python")]
227pub fn register_python_data_class(type_name: &str, data_class: &Bound<'_, PyAny>) {
228 python_data_classes().insert(type_name.to_string(), data_class.clone().unbind());
229}
230
231#[cfg(feature = "python")]
232#[must_use]
233pub fn get_python_data_class(py: Python<'_>, type_name: &str) -> Option<Py<PyAny>> {
234 python_data_classes()
235 .get(type_name)
236 .map(|entry| entry.value().clone_ref(py))
237}
238
239#[cfg(feature = "python")]
245pub fn reconstruct_python_custom_data(
246 py: Python<'_>,
247 type_name: &str,
248 json: &str,
249) -> PyResult<Py<PyAny>> {
250 let data_class = get_python_data_class(py, type_name).ok_or_else(|| {
251 nautilus_core::python::to_pyruntime_err(format!(
252 "No registered Python class for custom data type `{type_name}`"
253 ))
254 })?;
255 let json_module = py.import("json")?;
256 let payload = json_module.call_method1("loads", (json,))?;
257 data_class
258 .bind(py)
259 .call_method1("from_json", (payload,))
260 .map(Bound::unbind)
261}
262
263#[cfg(feature = "python")]
272pub fn clone_pyclass_to_pyobject<T>(value: &T, py: Python<'_>) -> PyResult<Py<PyAny>>
273where
274 T: Clone,
275 for<'py> T: pyo3::IntoPyObject<'py, Error = pyo3::PyErr>,
276{
277 value.clone().into_py_any(py)
278}
279
280pub trait CustomDataTrait: HasTsInit + Send + Sync + Debug {
282 fn type_name(&self) -> &'static str;
284
285 fn as_any(&self) -> &dyn Any;
287
288 fn ts_event(&self) -> UnixNanos;
290
291 fn to_json(&self) -> anyhow::Result<String>;
296
297 fn to_json_py(&self) -> anyhow::Result<String> {
303 self.to_json()
304 }
305
306 fn clone_arc(&self) -> Arc<dyn CustomDataTrait>;
308
309 fn eq_arc(&self, other: &dyn CustomDataTrait) -> bool;
311
312 #[cfg(feature = "python")]
317 fn to_pyobject(&self, _py: Python<'_>) -> PyResult<Py<PyAny>> {
318 Err(nautilus_core::python::to_pytype_err(format!(
319 "to_pyobject not implemented for {}",
320 self.type_name()
321 )))
322 }
323
324 #[must_use]
326 fn type_name_static() -> &'static str
327 where
328 Self: Sized,
329 {
330 std::any::type_name::<Self>()
331 }
332
333 fn from_json(_value: serde_json::Value) -> anyhow::Result<Arc<dyn CustomDataTrait>>
338 where
339 Self: Sized,
340 {
341 anyhow::bail!(
342 "from_json not implemented for {}",
343 std::any::type_name::<Self>()
344 )
345 }
346}
347
348pub fn register_custom_data_json<T: CustomDataTrait + Sized>() -> anyhow::Result<()> {
354 let type_name = T::type_name_static();
355 register_json_deserializer(type_name, Box::new(|value| T::from_json(value)))
356}
357
358#[allow(
363 clippy::missing_errors_doc,
364 reason = "idempotent registration only inserts into a DashMap"
365)]
366pub fn ensure_custom_data_json_registered<T: CustomDataTrait + Sized>() -> anyhow::Result<()> {
367 let type_name = T::type_name_static();
368 ensure_json_deserializer_registered(type_name, Box::new(|value| T::from_json(value)))
369}
370
371#[cfg_attr(
377 feature = "python",
378 pyclass(module = "nautilus_trader.model", name = "CustomData", from_py_object)
379)]
380#[cfg_attr(
381 feature = "python",
382 pyo3_stub_gen::derive::gen_stub_pyclass(module = "nautilus_trader.model")
383)]
384#[derive(Clone, Debug)]
385pub struct CustomData {
386 pub data: Arc<dyn CustomDataTrait>,
388 pub data_type: DataType,
390}
391
392impl CustomData {
393 pub fn from_arc(arc: Arc<dyn CustomDataTrait>) -> Self {
396 let data_type = DataType::new(arc.type_name(), None, None);
397 Self {
398 data: arc,
399 data_type,
400 }
401 }
402
403 pub fn new(data: Arc<dyn CustomDataTrait>, data_type: DataType) -> Self {
408 Self { data, data_type }
409 }
410}
411
412impl PartialEq for CustomData {
413 fn eq(&self, other: &Self) -> bool {
414 self.data.eq_arc(other.data.as_ref()) && self.data_type == other.data_type
415 }
416}
417
418impl HasTsInit for CustomData {
419 fn ts_init(&self) -> UnixNanos {
420 self.data.ts_init()
421 }
422}
423
424pub(crate) fn parse_custom_data_from_json_bytes(
425 bytes: &[u8],
426) -> Result<CustomData, serde_json::Error> {
427 let data: Data = serde_json::from_slice(bytes)?;
428 match data {
429 Data::Custom(custom) => Ok(custom),
430 _ => Err(serde_json::Error::io(std::io::Error::new(
431 std::io::ErrorKind::InvalidData,
432 "JSON does not represent CustomData",
433 ))),
434 }
435}
436
437impl CustomData {
438 pub fn from_json_bytes(bytes: &[u8]) -> Result<Self, serde_json::Error> {
444 parse_custom_data_from_json_bytes(bytes)
445 }
446}
447
448struct CustomDataEnvelope {
451 type_name: String,
452 data_type: serde_json::Value,
453 payload: serde_json::Value,
454}
455
456impl Serialize for CustomDataEnvelope {
457 fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
458 where
459 S: Serializer,
460 {
461 use serde::ser::SerializeStruct;
462 let mut state = serializer.serialize_struct("CustomDataEnvelope", 3)?;
463 state.serialize_field("type", &self.type_name)?;
464 state.serialize_field("data_type", &self.data_type)?;
465 state.serialize_field("payload", &self.payload)?;
466 state.end()
467 }
468}
469
470impl CustomData {
471 fn to_envelope_json_value(&self) -> Result<serde_json::Value, serde_json::Error> {
472 let json = self.data.to_json().map_err(|e| {
473 serde_json::Error::io(std::io::Error::new(std::io::ErrorKind::InvalidData, e))
474 })?;
475 let payload: serde_json::Value = serde_json::from_str(&json)?;
476 let metadata_value = self.data_type.metadata().map_or(
477 serde_json::Value::Object(serde_json::Map::new()),
478 |m| {
479 serde_json::to_value(m).unwrap_or(serde_json::Value::Object(serde_json::Map::new()))
480 },
481 );
482 let mut data_type_obj = serde_json::Map::new();
483 data_type_obj.insert(
484 "type_name".to_string(),
485 serde_json::Value::String(self.data_type.type_name().to_string()),
486 );
487 data_type_obj.insert("metadata".to_string(), metadata_value);
488
489 if let Some(id) = self.data_type.identifier() {
490 data_type_obj.insert(
491 "identifier".to_string(),
492 serde_json::Value::String(id.to_string()),
493 );
494 }
495
496 let envelope = CustomDataEnvelope {
497 type_name: self.data.type_name().to_string(),
498 data_type: serde_json::Value::Object(data_type_obj),
499 payload,
500 };
501 serde_json::to_value(envelope)
502 }
503}
504
505impl Serialize for CustomData {
506 fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
507 where
508 S: Serializer,
509 {
510 let value = self
511 .to_envelope_json_value()
512 .map_err(serde::ser::Error::custom)?;
513 value.serialize(serializer)
514 }
515}
516
517#[cfg(test)]
518mod tests {
519 use nautilus_core::{Params, UnixNanos};
520 use rstest::rstest;
521 use serde::Deserialize;
522 use serde_json::json;
523
524 use super::*;
525 use crate::{data::HasTsInit, identifiers::InstrumentId};
526
527 #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
528 struct TestCustomData {
529 ts_init: UnixNanos,
530 instrument_id: InstrumentId,
531 payload: serde_json::Value,
532 }
533
534 impl HasTsInit for TestCustomData {
535 fn ts_init(&self) -> UnixNanos {
536 self.ts_init
537 }
538 }
539
540 impl CustomDataTrait for TestCustomData {
541 fn type_name(&self) -> &'static str {
542 "TestCustomData"
543 }
544 fn as_any(&self) -> &dyn Any {
545 self
546 }
547 fn ts_event(&self) -> UnixNanos {
548 self.ts_init
549 }
550 fn to_json(&self) -> anyhow::Result<String> {
551 Ok(serde_json::to_string(self)?)
552 }
553 fn clone_arc(&self) -> Arc<dyn CustomDataTrait> {
554 Arc::new(self.clone())
555 }
556 fn eq_arc(&self, other: &dyn CustomDataTrait) -> bool {
557 if let Some(other) = other.as_any().downcast_ref::<Self>() {
558 self == other
559 } else {
560 false
561 }
562 }
563
564 fn type_name_static() -> &'static str {
565 "TestCustomData"
566 }
567
568 fn from_json(value: serde_json::Value) -> anyhow::Result<Arc<dyn CustomDataTrait>> {
569 let parsed: Self = serde_json::from_value(value)?;
570 Ok(Arc::new(parsed))
571 }
572 }
573
574 #[rstest]
575 fn test_custom_data_roundtrip(
576 #[values(false, true)] msgpack: bool,
577 #[values(None, Some(Params::empty()), Some(serde_json::from_value(json!({"key": [1, {"nested": true}]})).unwrap()))]
578 metadata: Option<Params>,
579 #[values(None, Some("TEST.SIM".to_string()))] identifier: Option<String>,
580 ) {
581 ensure_custom_data_json_registered::<TestCustomData>()
582 .expect("TestCustomData must register for JSON roundtrip test");
583
584 let instrument_id = InstrumentId::from("TEST.SIM");
585
586 let inner = TestCustomData {
587 ts_init: UnixNanos::from(100),
588 instrument_id,
589 payload: json!({"nested": [null, {"min": i64::MIN, "max": u64::MAX, "price": "123.4500"}], "empty": [], "flag": false}),
590 };
591
592 let data_type = DataType::new("TestCustomData", metadata, identifier);
593 let original = CustomData::new(Arc::new(inner), data_type);
594 let roundtripped = if msgpack {
595 let bytes = rmp_serde::to_vec_named(&original).unwrap();
596
597 let Data::Custom(value) = rmp_serde::from_slice(&bytes).unwrap() else {
598 panic!("custom payload must retain its variant");
599 };
600
601 value
602 } else {
603 let bytes = serde_json::to_vec(&original).unwrap();
604 CustomData::from_json_bytes(&bytes).unwrap()
605 };
606
607 assert_eq!(
608 roundtripped.data_type.type_name(),
609 original.data_type.type_name()
610 );
611 assert_eq!(
612 roundtripped.data_type.metadata(),
613 original
614 .data_type
615 .metadata()
616 .filter(|value| !value.is_empty())
617 );
618 assert_eq!(
619 roundtripped.data_type.identifier(),
620 original.data_type.identifier()
621 );
622 let orig_inner = original
623 .data
624 .as_any()
625 .downcast_ref::<TestCustomData>()
626 .unwrap();
627 let rt_inner = roundtripped
628 .data
629 .as_any()
630 .downcast_ref::<TestCustomData>()
631 .unwrap();
632 assert_eq!(orig_inner, rt_inner);
633 }
634
635 #[rstest]
636 fn test_custom_data_wrapper() {
637 let instrument_id = InstrumentId::from("TEST.SIM");
638 let data = TestCustomData {
639 ts_init: UnixNanos::from(100),
640 instrument_id,
641 payload: json!({"value": 42}),
642 };
643 let data_type = DataType::new("TestCustomData", None, Some(instrument_id.to_string()));
644 let custom_data = CustomData::new(Arc::new(data), data_type);
645
646 assert_eq!(custom_data.data.ts_init(), UnixNanos::from(100));
647 assert_eq!(Data::Custom(custom_data).instrument_id(), instrument_id);
648 }
649}