1#[cfg(feature = "python")]
17use std::collections::HashSet;
18#[cfg(feature = "python")]
19use std::sync::RwLock;
20use std::{any::Any, fmt::Debug, sync::Arc};
21
22use nautilus_core::UnixNanos;
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 if let Ok(guard) = set.read()
39 && guard.contains(name.as_str())
40 {
41 return guard.get(name.as_str()).copied().unwrap();
42 }
43
44 if let Ok(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 } else {
52 log::warn!("intern_type_name_static: RwLock poisoned, interning skipped for type name");
53 Box::leak(name.into_boxed_str())
54 }
55}
56
57#[cfg(feature = "python")]
64pub struct PythonCustomDataWrapper {
65 py_object: Py<PyAny>,
67 cached_ts_event: UnixNanos,
69 cached_ts_init: UnixNanos,
71 cached_type_name: String,
73 cached_type_name_static: &'static str,
75}
76
77#[cfg(feature = "python")]
78impl PythonCustomDataWrapper {
79 pub fn new(_py: Python<'_>, py_object: &Bound<'_, PyAny>) -> PyResult<Self> {
86 let ts_event: u64 = py_object.getattr("ts_event")?.extract()?;
88 let ts_event = UnixNanos::from(ts_event);
89
90 let ts_init: u64 = py_object.getattr("ts_init")?.extract()?;
92 let ts_init = UnixNanos::from(ts_init);
93
94 let data_class = py_object.get_type();
96 let type_name: String = if data_class.hasattr("type_name_static")? {
97 data_class.call_method0("type_name_static")?.extract()?
98 } else {
99 data_class.getattr("__name__")?.extract()?
100 };
101
102 let type_name_static: &'static str = intern_type_name_static(type_name.clone());
104
105 Ok(Self {
106 py_object: py_object.clone().unbind(),
107 cached_ts_event: ts_event,
108 cached_ts_init: ts_init,
109 cached_type_name: type_name,
110 cached_type_name_static: type_name_static,
111 })
112 }
113
114 #[must_use]
116 pub fn py_object(&self) -> &Py<PyAny> {
117 &self.py_object
118 }
119
120 #[must_use]
122 pub fn get_type_name(&self) -> &str {
123 &self.cached_type_name
124 }
125}
126
127#[cfg(feature = "python")]
128impl Clone for PythonCustomDataWrapper {
129 fn clone(&self) -> Self {
130 Python::attach(|py| Self {
131 py_object: self.py_object.clone_ref(py),
132 cached_ts_event: self.cached_ts_event,
133 cached_ts_init: self.cached_ts_init,
134 cached_type_name: self.cached_type_name.clone(),
135 cached_type_name_static: self.cached_type_name_static,
136 })
137 }
138}
139
140#[cfg(feature = "python")]
141impl Debug for PythonCustomDataWrapper {
142 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
143 f.debug_struct(stringify!(PythonCustomDataWrapper))
144 .field("py_object", &self.py_object)
145 .field("type_name", &self.cached_type_name)
146 .field("type_name_static", &self.cached_type_name_static)
147 .field("ts_event", &self.cached_ts_event)
148 .field("ts_init", &self.cached_ts_init)
149 .finish()
150 }
151}
152
153#[cfg(feature = "python")]
154impl HasTsInit for PythonCustomDataWrapper {
155 fn ts_init(&self) -> UnixNanos {
156 self.cached_ts_init
157 }
158}
159
160#[cfg(feature = "python")]
161impl CustomDataTrait for PythonCustomDataWrapper {
162 fn type_name(&self) -> &'static str {
163 self.cached_type_name_static
164 }
165
166 fn as_any(&self) -> &dyn Any {
167 self
168 }
169
170 fn ts_event(&self) -> UnixNanos {
171 self.cached_ts_event
172 }
173
174 fn to_json(&self) -> anyhow::Result<String> {
175 Python::attach(|py| {
176 let obj = self.py_object.bind(py);
177 if obj.hasattr("to_json")? {
179 let json_str: String = obj.call_method0("to_json")?.extract()?;
180 Ok(json_str)
181 } else {
182 let json_module = py.import("json")?;
184 let dict = if obj.hasattr("__dict__")? {
186 obj.getattr("__dict__")?
187 } else {
188 anyhow::bail!("Python object has no to_json() method or __dict__ attribute");
189 };
190 let json_str: String = json_module.call_method1("dumps", (dict,))?.extract()?;
191 Ok(json_str)
192 }
193 })
194 }
195
196 fn clone_arc(&self) -> Arc<dyn CustomDataTrait> {
197 Arc::new(self.clone())
198 }
199
200 fn eq_arc(&self, other: &dyn CustomDataTrait) -> bool {
201 if let Some(other_wrapper) = other.as_any().downcast_ref::<Self>() {
204 Python::attach(|py| {
205 let a = self.py_object.bind(py);
206 let b = other_wrapper.py_object.bind(py);
207 if a.is(b) {
208 return true;
209 }
210 a.eq(b).unwrap_or(false)
211 })
212 } else {
213 false
214 }
215 }
216
217 fn to_pyobject(&self, py: Python<'_>) -> PyResult<Py<PyAny>> {
218 Ok(self.py_object.clone_ref(py))
220 }
221}
222
223#[cfg(feature = "python")]
224fn python_data_classes() -> &'static dashmap::DashMap<String, Py<PyAny>> {
225 static PYTHON_DATA_CLASSES: std::sync::OnceLock<dashmap::DashMap<String, Py<PyAny>>> =
226 std::sync::OnceLock::new();
227 PYTHON_DATA_CLASSES.get_or_init(dashmap::DashMap::new)
228}
229
230#[cfg(feature = "python")]
231pub fn register_python_data_class(type_name: &str, data_class: &Bound<'_, PyAny>) {
232 python_data_classes().insert(type_name.to_string(), data_class.clone().unbind());
233}
234
235#[cfg(feature = "python")]
236#[must_use]
237pub fn get_python_data_class(py: Python<'_>, type_name: &str) -> Option<Py<PyAny>> {
238 python_data_classes()
239 .get(type_name)
240 .map(|entry| entry.value().clone_ref(py))
241}
242
243#[cfg(feature = "python")]
249pub fn reconstruct_python_custom_data(
250 py: Python<'_>,
251 type_name: &str,
252 json: &str,
253) -> PyResult<Py<PyAny>> {
254 let data_class = get_python_data_class(py, type_name).ok_or_else(|| {
255 nautilus_core::python::to_pyruntime_err(format!(
256 "No registered Python class for custom data type `{type_name}`"
257 ))
258 })?;
259 let json_module = py.import("json")?;
260 let payload = json_module.call_method1("loads", (json,))?;
261 data_class
262 .bind(py)
263 .call_method1("from_json", (payload,))
264 .map(Bound::unbind)
265}
266
267#[cfg(feature = "python")]
276pub fn clone_pyclass_to_pyobject<T>(value: &T, py: Python<'_>) -> PyResult<Py<PyAny>>
277where
278 T: Clone,
279 for<'py> T: pyo3::IntoPyObject<'py, Error = pyo3::PyErr>,
280{
281 value.clone().into_py_any(py)
282}
283
284pub trait CustomDataTrait: HasTsInit + Send + Sync + Debug {
286 fn type_name(&self) -> &'static str;
288
289 fn as_any(&self) -> &dyn Any;
291
292 fn ts_event(&self) -> UnixNanos;
294
295 fn to_json(&self) -> anyhow::Result<String>;
300
301 fn to_json_py(&self) -> anyhow::Result<String> {
307 self.to_json()
308 }
309
310 fn clone_arc(&self) -> Arc<dyn CustomDataTrait>;
312
313 fn eq_arc(&self, other: &dyn CustomDataTrait) -> bool;
315
316 #[cfg(feature = "python")]
321 fn to_pyobject(&self, _py: Python<'_>) -> PyResult<Py<PyAny>> {
322 Err(nautilus_core::python::to_pytype_err(format!(
323 "to_pyobject not implemented for {}",
324 self.type_name()
325 )))
326 }
327
328 #[must_use]
330 fn type_name_static() -> &'static str
331 where
332 Self: Sized,
333 {
334 std::any::type_name::<Self>()
335 }
336
337 fn from_json(_value: serde_json::Value) -> anyhow::Result<Arc<dyn CustomDataTrait>>
342 where
343 Self: Sized,
344 {
345 anyhow::bail!(
346 "from_json not implemented for {}",
347 std::any::type_name::<Self>()
348 )
349 }
350}
351
352pub fn register_custom_data_json<T: CustomDataTrait + Sized>() -> anyhow::Result<()> {
358 let type_name = T::type_name_static();
359 register_json_deserializer(type_name, Box::new(|value| T::from_json(value)))
360}
361
362pub fn ensure_custom_data_json_registered<T: CustomDataTrait + Sized>() -> anyhow::Result<()> {
368 let type_name = T::type_name_static();
369 ensure_json_deserializer_registered(type_name, Box::new(|value| T::from_json(value)))
370}
371
372#[cfg_attr(
378 feature = "python",
379 pyclass(module = "nautilus_trader.model", name = "CustomData", from_py_object)
380)]
381#[cfg_attr(
382 feature = "python",
383 pyo3_stub_gen::derive::gen_stub_pyclass(module = "nautilus_trader.model")
384)]
385#[derive(Clone, Debug)]
386pub struct CustomData {
387 pub data: Arc<dyn CustomDataTrait>,
389 pub data_type: DataType,
391}
392
393impl CustomData {
394 pub fn from_arc(arc: Arc<dyn CustomDataTrait>) -> Self {
397 let data_type = DataType::new(arc.type_name(), None, None);
398 Self {
399 data: arc,
400 data_type,
401 }
402 }
403
404 pub fn new(data: Arc<dyn CustomDataTrait>, data_type: DataType) -> Self {
409 Self { data, data_type }
410 }
411}
412
413impl PartialEq for CustomData {
414 fn eq(&self, other: &Self) -> bool {
415 self.data.eq_arc(other.data.as_ref()) && self.data_type == other.data_type
416 }
417}
418
419impl HasTsInit for CustomData {
420 fn ts_init(&self) -> UnixNanos {
421 self.data.ts_init()
422 }
423}
424
425pub(crate) fn parse_custom_data_from_json_bytes(
426 bytes: &[u8],
427) -> Result<CustomData, serde_json::Error> {
428 let data: Data = serde_json::from_slice(bytes)?;
429 match data {
430 Data::Custom(custom) => Ok(custom),
431 _ => Err(serde_json::Error::io(std::io::Error::new(
432 std::io::ErrorKind::InvalidData,
433 "JSON does not represent CustomData",
434 ))),
435 }
436}
437
438impl CustomData {
439 pub fn from_json_bytes(bytes: &[u8]) -> Result<Self, serde_json::Error> {
445 parse_custom_data_from_json_bytes(bytes)
446 }
447}
448
449struct CustomDataEnvelope {
452 type_name: String,
453 data_type: serde_json::Value,
454 payload: serde_json::Value,
455}
456
457impl Serialize for CustomDataEnvelope {
458 fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
459 where
460 S: Serializer,
461 {
462 use serde::ser::SerializeStruct;
463 let mut state = serializer.serialize_struct("CustomDataEnvelope", 3)?;
464 state.serialize_field("type", &self.type_name)?;
465 state.serialize_field("data_type", &self.data_type)?;
466 state.serialize_field("payload", &self.payload)?;
467 state.end()
468 }
469}
470
471impl CustomData {
472 fn to_envelope_json_value(&self) -> Result<serde_json::Value, serde_json::Error> {
473 let json = self.data.to_json().map_err(|e| {
474 serde_json::Error::io(std::io::Error::new(std::io::ErrorKind::InvalidData, e))
475 })?;
476 let payload: serde_json::Value = serde_json::from_str(&json)?;
477 let metadata_value = self.data_type.metadata().map_or(
478 serde_json::Value::Object(serde_json::Map::new()),
479 |m| {
480 serde_json::to_value(m).unwrap_or(serde_json::Value::Object(serde_json::Map::new()))
481 },
482 );
483 let mut data_type_obj = serde_json::Map::new();
484 data_type_obj.insert(
485 "type_name".to_string(),
486 serde_json::Value::String(self.data_type.type_name().to_string()),
487 );
488 data_type_obj.insert("metadata".to_string(), metadata_value);
489
490 if let Some(id) = self.data_type.identifier() {
491 data_type_obj.insert(
492 "identifier".to_string(),
493 serde_json::Value::String(id.to_string()),
494 );
495 }
496
497 let envelope = CustomDataEnvelope {
498 type_name: self.data.type_name().to_string(),
499 data_type: serde_json::Value::Object(data_type_obj),
500 payload,
501 };
502 serde_json::to_value(envelope)
503 }
504}
505
506impl Serialize for CustomData {
507 fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
508 where
509 S: Serializer,
510 {
511 let value = self
512 .to_envelope_json_value()
513 .map_err(serde::ser::Error::custom)?;
514 value.serialize(serializer)
515 }
516}
517
518#[cfg(test)]
519mod tests {
520 use nautilus_core::{Params, UnixNanos};
521 use rstest::rstest;
522 use serde::Deserialize;
523 use serde_json::json;
524
525 use super::*;
526 use crate::{data::HasTsInit, identifiers::InstrumentId};
527
528 #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
529 struct TestCustomData {
530 ts_init: UnixNanos,
531 instrument_id: InstrumentId,
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_json_roundtrip() {
576 register_custom_data_json::<TestCustomData>()
577 .expect("TestCustomData must register for JSON roundtrip test");
578
579 let instrument_id = InstrumentId::from("TEST.SIM");
580 let metadata = Some(
581 serde_json::from_value::<Params>(json!({"key1": "value1", "key2": "value2"})).unwrap(),
582 );
583 let inner = TestCustomData {
584 ts_init: UnixNanos::from(100),
585 instrument_id,
586 };
587 let data_type = DataType::new("TestCustomData", metadata, Some(instrument_id.to_string()));
588 let original = CustomData::new(Arc::new(inner), data_type);
589
590 let json_bytes = serde_json::to_vec(&original).unwrap();
591 let roundtripped = CustomData::from_json_bytes(&json_bytes).unwrap();
592
593 assert_eq!(
594 roundtripped.data_type.type_name(),
595 original.data_type.type_name()
596 );
597 assert_eq!(
598 roundtripped.data_type.metadata(),
599 original.data_type.metadata()
600 );
601 assert_eq!(
602 roundtripped.data_type.identifier(),
603 original.data_type.identifier()
604 );
605 let orig_inner = original
606 .data
607 .as_any()
608 .downcast_ref::<TestCustomData>()
609 .unwrap();
610 let rt_inner = roundtripped
611 .data
612 .as_any()
613 .downcast_ref::<TestCustomData>()
614 .unwrap();
615 assert_eq!(orig_inner, rt_inner);
616 }
617
618 #[rstest]
619 fn test_custom_data_wrapper() {
620 let instrument_id = InstrumentId::from("TEST.SIM");
621 let data = TestCustomData {
622 ts_init: UnixNanos::from(100),
623 instrument_id,
624 };
625 let data_type = DataType::new("TestCustomData", None, Some(instrument_id.to_string()));
626 let custom_data = CustomData::new(Arc::new(data), data_type);
627
628 assert_eq!(custom_data.data.ts_init(), UnixNanos::from(100));
629 assert_eq!(Data::Custom(custom_data).instrument_id(), instrument_id);
630 }
631}