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 guard.contains(name.as_str()) {
40 return guard.get(name.as_str()).copied().unwrap();
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
358pub fn ensure_custom_data_json_registered<T: CustomDataTrait + Sized>() -> anyhow::Result<()> {
364 let type_name = T::type_name_static();
365 ensure_json_deserializer_registered(type_name, Box::new(|value| T::from_json(value)))
366}
367
368#[cfg_attr(
374 feature = "python",
375 pyclass(module = "nautilus_trader.model", name = "CustomData", from_py_object)
376)]
377#[cfg_attr(
378 feature = "python",
379 pyo3_stub_gen::derive::gen_stub_pyclass(module = "nautilus_trader.model")
380)]
381#[derive(Clone, Debug)]
382pub struct CustomData {
383 pub data: Arc<dyn CustomDataTrait>,
385 pub data_type: DataType,
387}
388
389impl CustomData {
390 pub fn from_arc(arc: Arc<dyn CustomDataTrait>) -> Self {
393 let data_type = DataType::new(arc.type_name(), None, None);
394 Self {
395 data: arc,
396 data_type,
397 }
398 }
399
400 pub fn new(data: Arc<dyn CustomDataTrait>, data_type: DataType) -> Self {
405 Self { data, data_type }
406 }
407}
408
409impl PartialEq for CustomData {
410 fn eq(&self, other: &Self) -> bool {
411 self.data.eq_arc(other.data.as_ref()) && self.data_type == other.data_type
412 }
413}
414
415impl HasTsInit for CustomData {
416 fn ts_init(&self) -> UnixNanos {
417 self.data.ts_init()
418 }
419}
420
421pub(crate) fn parse_custom_data_from_json_bytes(
422 bytes: &[u8],
423) -> Result<CustomData, serde_json::Error> {
424 let data: Data = serde_json::from_slice(bytes)?;
425 match data {
426 Data::Custom(custom) => Ok(custom),
427 _ => Err(serde_json::Error::io(std::io::Error::new(
428 std::io::ErrorKind::InvalidData,
429 "JSON does not represent CustomData",
430 ))),
431 }
432}
433
434impl CustomData {
435 pub fn from_json_bytes(bytes: &[u8]) -> Result<Self, serde_json::Error> {
441 parse_custom_data_from_json_bytes(bytes)
442 }
443}
444
445struct CustomDataEnvelope {
448 type_name: String,
449 data_type: serde_json::Value,
450 payload: serde_json::Value,
451}
452
453impl Serialize for CustomDataEnvelope {
454 fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
455 where
456 S: Serializer,
457 {
458 use serde::ser::SerializeStruct;
459 let mut state = serializer.serialize_struct("CustomDataEnvelope", 3)?;
460 state.serialize_field("type", &self.type_name)?;
461 state.serialize_field("data_type", &self.data_type)?;
462 state.serialize_field("payload", &self.payload)?;
463 state.end()
464 }
465}
466
467impl CustomData {
468 fn to_envelope_json_value(&self) -> Result<serde_json::Value, serde_json::Error> {
469 let json = self.data.to_json().map_err(|e| {
470 serde_json::Error::io(std::io::Error::new(std::io::ErrorKind::InvalidData, e))
471 })?;
472 let payload: serde_json::Value = serde_json::from_str(&json)?;
473 let metadata_value = self.data_type.metadata().map_or(
474 serde_json::Value::Object(serde_json::Map::new()),
475 |m| {
476 serde_json::to_value(m).unwrap_or(serde_json::Value::Object(serde_json::Map::new()))
477 },
478 );
479 let mut data_type_obj = serde_json::Map::new();
480 data_type_obj.insert(
481 "type_name".to_string(),
482 serde_json::Value::String(self.data_type.type_name().to_string()),
483 );
484 data_type_obj.insert("metadata".to_string(), metadata_value);
485
486 if let Some(id) = self.data_type.identifier() {
487 data_type_obj.insert(
488 "identifier".to_string(),
489 serde_json::Value::String(id.to_string()),
490 );
491 }
492
493 let envelope = CustomDataEnvelope {
494 type_name: self.data.type_name().to_string(),
495 data_type: serde_json::Value::Object(data_type_obj),
496 payload,
497 };
498 serde_json::to_value(envelope)
499 }
500}
501
502impl Serialize for CustomData {
503 fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
504 where
505 S: Serializer,
506 {
507 let value = self
508 .to_envelope_json_value()
509 .map_err(serde::ser::Error::custom)?;
510 value.serialize(serializer)
511 }
512}
513
514#[cfg(test)]
515mod tests {
516 use nautilus_core::{Params, UnixNanos};
517 use rstest::rstest;
518 use serde::Deserialize;
519 use serde_json::json;
520
521 use super::*;
522 use crate::{data::HasTsInit, identifiers::InstrumentId};
523
524 #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
525 struct TestCustomData {
526 ts_init: UnixNanos,
527 instrument_id: InstrumentId,
528 }
529
530 impl HasTsInit for TestCustomData {
531 fn ts_init(&self) -> UnixNanos {
532 self.ts_init
533 }
534 }
535
536 impl CustomDataTrait for TestCustomData {
537 fn type_name(&self) -> &'static str {
538 "TestCustomData"
539 }
540 fn as_any(&self) -> &dyn Any {
541 self
542 }
543 fn ts_event(&self) -> UnixNanos {
544 self.ts_init
545 }
546 fn to_json(&self) -> anyhow::Result<String> {
547 Ok(serde_json::to_string(self)?)
548 }
549 fn clone_arc(&self) -> Arc<dyn CustomDataTrait> {
550 Arc::new(self.clone())
551 }
552 fn eq_arc(&self, other: &dyn CustomDataTrait) -> bool {
553 if let Some(other) = other.as_any().downcast_ref::<Self>() {
554 self == other
555 } else {
556 false
557 }
558 }
559
560 fn type_name_static() -> &'static str {
561 "TestCustomData"
562 }
563
564 fn from_json(value: serde_json::Value) -> anyhow::Result<Arc<dyn CustomDataTrait>> {
565 let parsed: Self = serde_json::from_value(value)?;
566 Ok(Arc::new(parsed))
567 }
568 }
569
570 #[rstest]
571 fn test_custom_data_json_roundtrip() {
572 register_custom_data_json::<TestCustomData>()
573 .expect("TestCustomData must register for JSON roundtrip test");
574
575 let instrument_id = InstrumentId::from("TEST.SIM");
576 let metadata = Some(
577 serde_json::from_value::<Params>(json!({"key1": "value1", "key2": "value2"})).unwrap(),
578 );
579 let inner = TestCustomData {
580 ts_init: UnixNanos::from(100),
581 instrument_id,
582 };
583 let data_type = DataType::new("TestCustomData", metadata, Some(instrument_id.to_string()));
584 let original = CustomData::new(Arc::new(inner), data_type);
585
586 let json_bytes = serde_json::to_vec(&original).unwrap();
587 let roundtripped = CustomData::from_json_bytes(&json_bytes).unwrap();
588
589 assert_eq!(
590 roundtripped.data_type.type_name(),
591 original.data_type.type_name()
592 );
593 assert_eq!(
594 roundtripped.data_type.metadata(),
595 original.data_type.metadata()
596 );
597 assert_eq!(
598 roundtripped.data_type.identifier(),
599 original.data_type.identifier()
600 );
601 let orig_inner = original
602 .data
603 .as_any()
604 .downcast_ref::<TestCustomData>()
605 .unwrap();
606 let rt_inner = roundtripped
607 .data
608 .as_any()
609 .downcast_ref::<TestCustomData>()
610 .unwrap();
611 assert_eq!(orig_inner, rt_inner);
612 }
613
614 #[rstest]
615 fn test_custom_data_wrapper() {
616 let instrument_id = InstrumentId::from("TEST.SIM");
617 let data = TestCustomData {
618 ts_init: UnixNanos::from(100),
619 instrument_id,
620 };
621 let data_type = DataType::new("TestCustomData", None, Some(instrument_id.to_string()));
622 let custom_data = CustomData::new(Arc::new(data), data_type);
623
624 assert_eq!(custom_data.data.ts_init(), UnixNanos::from(100));
625 assert_eq!(Data::Custom(custom_data).instrument_id(), instrument_id);
626 }
627}