1use nautilus_core::{
17 UUID4, UnixNanos,
18 python::{IntoPyObjectNautilusExt, to_pyvalue_err},
19};
20use nautilus_model::identifiers::{ClientId, TraderId, Venue};
21use pyo3::{basic::CompareOp, prelude::*};
22use ustr::Ustr;
23
24use crate::{
25 messages::system::{
26 QueueCondition, QueueState, QueueStateChanged, ReconnectSocket, SocketState,
27 SocketStateChanged, socket_endpoint,
28 },
29 runner::SystemChannel,
30};
31
32#[pymethods]
33#[pyo3_stub_gen::derive::gen_stub_pymethods]
34impl SystemChannel {
35 const fn __hash__(&self) -> isize {
36 *self as isize
37 }
38}
39
40#[pymethods]
41#[pyo3_stub_gen::derive::gen_stub_pymethods]
42impl QueueCondition {
43 const fn __hash__(&self) -> isize {
44 *self as isize
45 }
46}
47
48#[pymethods]
49#[pyo3_stub_gen::derive::gen_stub_pymethods]
50impl QueueState {
51 const fn __hash__(&self) -> isize {
52 *self as isize
53 }
54}
55
56#[pymethods]
57#[pyo3_stub_gen::derive::gen_stub_pymethods]
58impl SocketState {
59 const fn __hash__(&self) -> isize {
60 *self as isize
61 }
62}
63
64#[pymethods]
65#[pyo3_stub_gen::derive::gen_stub_pymethods]
66impl ReconnectSocket {
67 #[new]
69 fn py_new(
70 trader_id: TraderId,
71 client_id: ClientId,
72 endpoint: &str,
73 ts_init: u64,
74 ) -> PyResult<Self> {
75 Ok(Self::new(
76 trader_id,
77 client_id,
78 socket_endpoint(endpoint).map_err(to_pyvalue_err)?,
79 UnixNanos::from(ts_init),
80 ))
81 }
82
83 fn __richcmp__(&self, other: &Self, op: CompareOp, py: Python<'_>) -> Py<PyAny> {
84 match op {
85 CompareOp::Eq => self.eq(other).into_py_any_unwrap(py),
86 CompareOp::Ne => self.ne(other).into_py_any_unwrap(py),
87 _ => py.NotImplemented(),
88 }
89 }
90
91 fn __repr__(&self) -> String {
92 self.to_string()
93 }
94
95 #[getter]
96 const fn trader_id(&self) -> TraderId {
97 self.trader_id
98 }
99
100 #[getter]
101 const fn client_id(&self) -> ClientId {
102 self.client_id
103 }
104
105 #[getter]
106 fn endpoint(&self) -> &str {
107 self.endpoint.as_str()
108 }
109
110 #[getter]
111 const fn ts_init(&self) -> u64 {
112 self.ts_init.as_u64()
113 }
114}
115
116#[pymethods]
117#[pyo3_stub_gen::derive::gen_stub_pymethods]
118impl QueueStateChanged {
119 #[new]
121 #[expect(clippy::too_many_arguments)]
122 fn py_new(
123 trader_id: TraderId,
124 channel: SystemChannel,
125 condition: QueueCondition,
126 state: QueueState,
127 queue_depth: usize,
128 mean_dispatch_ns: u64,
129 event_id: UUID4,
130 ts_event: u64,
131 ts_init: u64,
132 ) -> Self {
133 Self::new(
134 trader_id,
135 channel,
136 condition,
137 state,
138 queue_depth,
139 mean_dispatch_ns,
140 event_id,
141 UnixNanos::from(ts_event),
142 UnixNanos::from(ts_init),
143 )
144 }
145
146 fn __richcmp__(&self, other: &Self, op: CompareOp, py: Python<'_>) -> Py<PyAny> {
147 match op {
148 CompareOp::Eq => self.eq(other).into_py_any_unwrap(py),
149 CompareOp::Ne => self.ne(other).into_py_any_unwrap(py),
150 _ => py.NotImplemented(),
151 }
152 }
153
154 fn __repr__(&self) -> String {
155 self.to_string()
156 }
157
158 #[getter]
159 #[pyo3(name = "trader_id")]
160 const fn py_trader_id(&self) -> TraderId {
161 self.trader_id
162 }
163
164 #[getter]
165 #[pyo3(name = "channel")]
166 const fn py_channel(&self) -> SystemChannel {
167 self.channel
168 }
169
170 #[getter]
171 #[pyo3(name = "condition")]
172 const fn py_condition(&self) -> QueueCondition {
173 self.condition
174 }
175
176 #[getter]
177 #[pyo3(name = "state")]
178 const fn py_state(&self) -> QueueState {
179 self.state
180 }
181
182 #[getter]
183 #[pyo3(name = "queue_depth")]
184 const fn py_queue_depth(&self) -> usize {
185 self.queue_depth
186 }
187
188 #[getter]
189 #[pyo3(name = "mean_dispatch_ns")]
190 const fn py_mean_dispatch_ns(&self) -> u64 {
191 self.mean_dispatch_ns
192 }
193
194 #[getter]
195 #[pyo3(name = "event_id")]
196 const fn py_event_id(&self) -> UUID4 {
197 self.event_id
198 }
199
200 #[getter]
201 #[pyo3(name = "ts_event")]
202 const fn py_ts_event(&self) -> u64 {
203 self.ts_event.as_u64()
204 }
205
206 #[getter]
207 #[pyo3(name = "ts_init")]
208 const fn py_ts_init(&self) -> u64 {
209 self.ts_init.as_u64()
210 }
211}
212
213#[pymethods]
214#[pyo3_stub_gen::derive::gen_stub_pymethods]
215impl SocketStateChanged {
216 #[new]
218 #[expect(clippy::too_many_arguments)]
219 fn py_new(
220 trader_id: TraderId,
221 client_id: ClientId,
222 venue: Option<Venue>,
223 endpoint: &str,
224 state: SocketState,
225 event_id: UUID4,
226 ts_event: u64,
227 ts_init: u64,
228 ) -> Self {
229 Self::new(
230 trader_id,
231 client_id,
232 venue,
233 Ustr::from(endpoint),
234 state,
235 event_id,
236 UnixNanos::from(ts_event),
237 UnixNanos::from(ts_init),
238 )
239 }
240
241 fn __richcmp__(&self, other: &Self, op: CompareOp, py: Python<'_>) -> Py<PyAny> {
242 match op {
243 CompareOp::Eq => self.eq(other).into_py_any_unwrap(py),
244 CompareOp::Ne => self.ne(other).into_py_any_unwrap(py),
245 _ => py.NotImplemented(),
246 }
247 }
248
249 fn __repr__(&self) -> String {
250 self.to_string()
251 }
252
253 #[getter]
254 #[pyo3(name = "trader_id")]
255 const fn py_trader_id(&self) -> TraderId {
256 self.trader_id
257 }
258
259 #[getter]
260 #[pyo3(name = "client_id")]
261 const fn py_client_id(&self) -> ClientId {
262 self.client_id
263 }
264
265 #[getter]
266 #[pyo3(name = "venue")]
267 const fn py_venue(&self) -> Option<Venue> {
268 self.venue
269 }
270
271 #[getter]
272 #[pyo3(name = "endpoint")]
273 fn py_endpoint(&self) -> &str {
274 self.endpoint.as_str()
275 }
276
277 #[getter]
278 #[pyo3(name = "state")]
279 const fn py_state(&self) -> SocketState {
280 self.state
281 }
282
283 #[getter]
284 #[pyo3(name = "event_id")]
285 const fn py_event_id(&self) -> UUID4 {
286 self.event_id
287 }
288
289 #[getter]
290 #[pyo3(name = "ts_event")]
291 const fn py_ts_event(&self) -> u64 {
292 self.ts_event.as_u64()
293 }
294
295 #[getter]
296 #[pyo3(name = "ts_init")]
297 const fn py_ts_init(&self) -> u64 {
298 self.ts_init.as_u64()
299 }
300}
301
302#[cfg(test)]
303mod tests {
304 use pyo3::exceptions::PyValueError;
305 use rstest::rstest;
306
307 use super::*;
308
309 #[rstest]
310 fn reconnect_socket_python_constructor_assigns_all_fields() {
311 let trader_id = TraderId::from("TRADER-001");
312 let client_id = ClientId::from("POLYMARKET");
313 let command =
314 ReconnectSocket::py_new(trader_id, client_id, "polymarket-market-streams", 11).unwrap();
315
316 assert_eq!(command.trader_id, trader_id);
317 assert_eq!(command.client_id, client_id);
318 assert_eq!(command.endpoint.as_str(), "polymarket-market-streams");
319 assert_eq!(command.ts_init, UnixNanos::from(11));
320 }
321
322 #[rstest]
323 fn reconnect_socket_python_constructor_rejects_raw_urls() {
324 pyo3::Python::initialize();
325 let trader_id = TraderId::from("TRADER-001");
326 let client_id = ClientId::from("POLYMARKET");
327
328 Python::attach(|py| {
329 let command_error = ReconnectSocket::py_new(
330 trader_id,
331 client_id,
332 "wss://user:secret@example.com/feed",
333 11,
334 )
335 .unwrap_err();
336
337 assert!(command_error.is_instance_of::<PyValueError>(py));
338 assert_eq!(
339 command_error.value(py).to_string(),
340 "Socket endpoint must contain only ASCII letters, digits, '.', '-', or '_'",
341 );
342 });
343 }
344}