nautilus_model/python/types/
quantity.rs1use std::{
17 collections::hash_map::DefaultHasher,
18 hash::{Hash, Hasher},
19 ops::Neg,
20};
21
22use nautilus_core::python::{
23 IntoPyObjectNautilusExt, get_pytype_name, to_pytype_err, to_pyvalue_err,
24};
25use pyo3::{
26 conversion::IntoPyObjectExt,
27 prelude::*,
28 pyclass::CompareOp,
29 types::{PyFloat, PyTuple},
30};
31use rust_decimal::{Decimal, RoundingStrategy};
32
33use crate::types::{Quantity, quantity::QuantityRaw};
34
35#[pymethods]
36impl Quantity {
37 #[new]
38 fn py_new(value: f64, precision: u8) -> PyResult<Self> {
39 Self::new_checked(value, precision).map_err(to_pyvalue_err)
40 }
41
42 fn __setstate__(&mut self, state: &Bound<'_, PyAny>) -> PyResult<()> {
43 let py_tuple: &Bound<'_, PyTuple> = state.downcast::<PyTuple>()?;
44 self.raw = py_tuple.get_item(0)?.extract::<QuantityRaw>()?;
45 self.precision = py_tuple.get_item(1)?.extract::<u8>()?;
46 Ok(())
47 }
48
49 fn __getstate__(&self, py: Python) -> PyResult<PyObject> {
50 (self.raw, self.precision).into_py_any(py)
51 }
52
53 fn __reduce__(&self, py: Python) -> PyResult<PyObject> {
54 let safe_constructor = py.get_type::<Self>().getattr("_safe_constructor")?;
55 let state = self.__getstate__(py)?;
56 (safe_constructor, PyTuple::empty(py), state).into_py_any(py)
57 }
58
59 #[staticmethod]
60 fn _safe_constructor() -> PyResult<Self> {
61 Ok(Self::zero(0)) }
63
64 fn __add__(&self, other: &Bound<'_, PyAny>, py: Python) -> PyResult<PyObject> {
65 if other.is_instance_of::<PyFloat>() {
66 let other_float: f64 = other.extract()?;
67 (self.as_f64() + other_float).into_py_any(py)
68 } else if let Ok(other_qty) = other.extract::<Self>() {
69 (self.as_decimal() + other_qty.as_decimal()).into_py_any(py)
70 } else if let Ok(other_dec) = other.extract::<Decimal>() {
71 (self.as_decimal() + other_dec).into_py_any(py)
72 } else {
73 let pytype_name = get_pytype_name(other)?;
74 Err(to_pytype_err(format!(
75 "Unsupported type for __add__, was `{pytype_name}`"
76 )))
77 }
78 }
79
80 fn __radd__(&self, other: &Bound<'_, PyAny>, py: Python) -> PyResult<PyObject> {
81 if other.is_instance_of::<PyFloat>() {
82 let other_float: f64 = other.extract()?;
83 (other_float + self.as_f64()).into_py_any(py)
84 } else if let Ok(other_qty) = other.extract::<Self>() {
85 (other_qty.as_decimal() + self.as_decimal()).into_py_any(py)
86 } else if let Ok(other_dec) = other.extract::<Decimal>() {
87 (other_dec + self.as_decimal()).into_py_any(py)
88 } else {
89 let pytype_name = get_pytype_name(other)?;
90 Err(to_pytype_err(format!(
91 "Unsupported type for __radd__, was `{pytype_name}`"
92 )))
93 }
94 }
95
96 fn __sub__(&self, other: &Bound<'_, PyAny>, py: Python) -> PyResult<PyObject> {
97 if other.is_instance_of::<PyFloat>() {
98 let other_float: f64 = other.extract()?;
99 (self.as_f64() - other_float).into_py_any(py)
100 } else if let Ok(other_qty) = other.extract::<Self>() {
101 (self.as_decimal() - other_qty.as_decimal()).into_py_any(py)
102 } else if let Ok(other_dec) = other.extract::<Decimal>() {
103 (self.as_decimal() - other_dec).into_py_any(py)
104 } else {
105 let pytype_name = get_pytype_name(other)?;
106 Err(to_pytype_err(format!(
107 "Unsupported type for __sub__, was `{pytype_name}`"
108 )))
109 }
110 }
111
112 fn __rsub__(&self, other: &Bound<'_, PyAny>, py: Python) -> PyResult<PyObject> {
113 if other.is_instance_of::<PyFloat>() {
114 let other_float: f64 = other.extract()?;
115 (other_float - self.as_f64()).into_py_any(py)
116 } else if let Ok(other_qty) = other.extract::<Self>() {
117 (other_qty.as_decimal() - self.as_decimal()).into_py_any(py)
118 } else if let Ok(other_dec) = other.extract::<Decimal>() {
119 (other_dec - self.as_decimal()).into_py_any(py)
120 } else {
121 let pytype_name = get_pytype_name(other)?;
122 Err(to_pytype_err(format!(
123 "Unsupported type for __rsub__, was `{pytype_name}`"
124 )))
125 }
126 }
127
128 fn __mul__(&self, other: &Bound<'_, PyAny>, py: Python) -> PyResult<PyObject> {
129 if other.is_instance_of::<PyFloat>() {
130 let other_float: f64 = other.extract()?;
131 (self.as_f64() * other_float).into_py_any(py)
132 } else if let Ok(other_qty) = other.extract::<Self>() {
133 (self.as_decimal() * other_qty.as_decimal()).into_py_any(py)
134 } else if let Ok(other_dec) = other.extract::<Decimal>() {
135 (self.as_decimal() * other_dec).into_py_any(py)
136 } else {
137 let pytype_name = get_pytype_name(other)?;
138 Err(to_pytype_err(format!(
139 "Unsupported type for __mul__, was `{pytype_name}`"
140 )))
141 }
142 }
143
144 fn __rmul__(&self, other: &Bound<'_, PyAny>, py: Python) -> PyResult<PyObject> {
145 if other.is_instance_of::<PyFloat>() {
146 let other_float: f64 = other.extract()?;
147 (other_float * self.as_f64()).into_py_any(py)
148 } else if let Ok(other_qty) = other.extract::<Self>() {
149 (other_qty.as_decimal() * self.as_decimal()).into_py_any(py)
150 } else if let Ok(other_dec) = other.extract::<Decimal>() {
151 (other_dec * self.as_decimal()).into_py_any(py)
152 } else {
153 let pytype_name = get_pytype_name(other)?;
154 Err(to_pytype_err(format!(
155 "Unsupported type for __rmul__, was `{pytype_name}`"
156 )))
157 }
158 }
159
160 fn __truediv__(&self, other: &Bound<'_, PyAny>, py: Python) -> PyResult<PyObject> {
161 if other.is_instance_of::<PyFloat>() {
162 let other_float: f64 = other.extract()?;
163 (self.as_f64() / other_float).into_py_any(py)
164 } else if let Ok(other_qty) = other.extract::<Self>() {
165 (self.as_decimal() / other_qty.as_decimal()).into_py_any(py)
166 } else if let Ok(other_dec) = other.extract::<Decimal>() {
167 (self.as_decimal() / other_dec).into_py_any(py)
168 } else {
169 let pytype_name = get_pytype_name(other)?;
170 Err(to_pytype_err(format!(
171 "Unsupported type for __truediv__, was `{pytype_name}`"
172 )))
173 }
174 }
175
176 fn __rtruediv__(&self, other: &Bound<'_, PyAny>, py: Python) -> PyResult<PyObject> {
177 if other.is_instance_of::<PyFloat>() {
178 let other_float: f64 = other.extract()?;
179 (other_float / self.as_f64()).into_py_any(py)
180 } else if let Ok(other_qty) = other.extract::<Self>() {
181 (other_qty.as_decimal() / self.as_decimal()).into_py_any(py)
182 } else if let Ok(other_dec) = other.extract::<Decimal>() {
183 (other_dec / self.as_decimal()).into_py_any(py)
184 } else {
185 let pytype_name = get_pytype_name(other)?;
186 Err(to_pytype_err(format!(
187 "Unsupported type for __rtruediv__, was `{pytype_name}`"
188 )))
189 }
190 }
191
192 fn __floordiv__(&self, other: &Bound<'_, PyAny>, py: Python) -> PyResult<PyObject> {
193 if other.is_instance_of::<PyFloat>() {
194 let other_float: f64 = other.extract()?;
195 (self.as_f64() / other_float).floor().into_py_any(py)
196 } else if let Ok(other_qty) = other.extract::<Self>() {
197 (self.as_decimal() / other_qty.as_decimal())
198 .floor()
199 .into_py_any(py)
200 } else if let Ok(other_dec) = other.extract::<Decimal>() {
201 (self.as_decimal() / other_dec).floor().into_py_any(py)
202 } else {
203 let pytype_name = get_pytype_name(other)?;
204 Err(to_pytype_err(format!(
205 "Unsupported type for __floordiv__, was `{pytype_name}`"
206 )))
207 }
208 }
209
210 fn __rfloordiv__(&self, other: &Bound<'_, PyAny>, py: Python) -> PyResult<PyObject> {
211 if other.is_instance_of::<PyFloat>() {
212 let other_float: f64 = other.extract()?;
213 (other_float / self.as_f64()).floor().into_py_any(py)
214 } else if let Ok(other_qty) = other.extract::<Self>() {
215 (other_qty.as_decimal() / self.as_decimal())
216 .floor()
217 .into_py_any(py)
218 } else if let Ok(other_dec) = other.extract::<Decimal>() {
219 (other_dec / self.as_decimal()).floor().into_py_any(py)
220 } else {
221 let pytype_name = get_pytype_name(other)?;
222 Err(to_pytype_err(format!(
223 "Unsupported type for __rfloordiv__, was `{pytype_name}`"
224 )))
225 }
226 }
227
228 fn __mod__(&self, other: &Bound<'_, PyAny>, py: Python) -> PyResult<PyObject> {
229 if other.is_instance_of::<PyFloat>() {
230 let other_float: f64 = other.extract()?;
231 (self.as_f64() % other_float).into_py_any(py)
232 } else if let Ok(other_qty) = other.extract::<Self>() {
233 (self.as_decimal() % other_qty.as_decimal()).into_py_any(py)
234 } else if let Ok(other_dec) = other.extract::<Decimal>() {
235 (self.as_decimal() % other_dec).into_py_any(py)
236 } else {
237 let pytype_name = get_pytype_name(other)?;
238 Err(to_pytype_err(format!(
239 "Unsupported type for __mod__, was `{pytype_name}`"
240 )))
241 }
242 }
243
244 fn __rmod__(&self, other: &Bound<'_, PyAny>, py: Python) -> PyResult<PyObject> {
245 if other.is_instance_of::<PyFloat>() {
246 let other_float: f64 = other.extract()?;
247 (other_float % self.as_f64()).into_py_any(py)
248 } else if let Ok(other_qty) = other.extract::<Self>() {
249 (other_qty.as_decimal() % self.as_decimal()).into_py_any(py)
250 } else if let Ok(other_dec) = other.extract::<Decimal>() {
251 (other_dec % self.as_decimal()).into_py_any(py)
252 } else {
253 let pytype_name = get_pytype_name(other)?;
254 Err(to_pytype_err(format!(
255 "Unsupported type for __rmod__, was `{pytype_name}`"
256 )))
257 }
258 }
259
260 fn __neg__(&self) -> Decimal {
261 self.as_decimal().neg()
262 }
263
264 fn __pos__(&self) -> Decimal {
265 let mut value = self.as_decimal();
266 value.set_sign_positive(true);
267 value
268 }
269
270 fn __abs__(&self) -> Decimal {
271 self.as_decimal().abs()
272 }
273
274 fn __int__(&self) -> u64 {
275 self.as_f64() as u64
276 }
277
278 fn __float__(&self) -> f64 {
279 self.as_f64()
280 }
281
282 #[pyo3(signature = (ndigits=None))]
283 fn __round__(&self, ndigits: Option<u32>) -> Decimal {
284 self.as_decimal()
285 .round_dp_with_strategy(ndigits.unwrap_or(0), RoundingStrategy::MidpointNearestEven)
286 }
287
288 fn __richcmp__(&self, other: PyObject, op: CompareOp, py: Python<'_>) -> Py<PyAny> {
289 if let Ok(other_qty) = other.extract::<Self>(py) {
290 match op {
291 CompareOp::Eq => self.eq(&other_qty).into_py_any_unwrap(py),
292 CompareOp::Ne => self.ne(&other_qty).into_py_any_unwrap(py),
293 CompareOp::Ge => self.ge(&other_qty).into_py_any_unwrap(py),
294 CompareOp::Gt => self.gt(&other_qty).into_py_any_unwrap(py),
295 CompareOp::Le => self.le(&other_qty).into_py_any_unwrap(py),
296 CompareOp::Lt => self.lt(&other_qty).into_py_any_unwrap(py),
297 }
298 } else if let Ok(other_dec) = other.extract::<Decimal>(py) {
299 match op {
300 CompareOp::Eq => (self.as_decimal() == other_dec).into_py_any_unwrap(py),
301 CompareOp::Ne => (self.as_decimal() != other_dec).into_py_any_unwrap(py),
302 CompareOp::Ge => (self.as_decimal() >= other_dec).into_py_any_unwrap(py),
303 CompareOp::Gt => (self.as_decimal() > other_dec).into_py_any_unwrap(py),
304 CompareOp::Le => (self.as_decimal() <= other_dec).into_py_any_unwrap(py),
305 CompareOp::Lt => (self.as_decimal() < other_dec).into_py_any_unwrap(py),
306 }
307 } else {
308 py.NotImplemented()
309 }
310 }
311
312 fn __hash__(&self) -> isize {
313 let mut h = DefaultHasher::new();
314 self.hash(&mut h);
315 h.finish() as isize
316 }
317
318 fn __repr__(&self) -> String {
319 format!("{self:?}")
320 }
321
322 fn __str__(&self) -> String {
323 self.to_string()
324 }
325
326 #[getter]
327 fn raw(&self) -> QuantityRaw {
328 self.raw
329 }
330
331 #[getter]
332 fn precision(&self) -> u8 {
333 self.precision
334 }
335
336 #[staticmethod]
337 #[pyo3(name = "from_raw")]
338 fn py_from_raw(raw: QuantityRaw, precision: u8) -> Self {
339 Self::from_raw(raw, precision)
340 }
341
342 #[staticmethod]
343 #[pyo3(name = "zero")]
344 #[pyo3(signature = (precision = 0))]
345 fn py_zero(precision: u8) -> PyResult<Self> {
346 Self::new_checked(0.0, precision).map_err(to_pyvalue_err)
347 }
348
349 #[staticmethod]
350 #[pyo3(name = "from_int")]
351 fn py_from_int(value: u64) -> PyResult<Self> {
352 Self::new_checked(value as f64, 0).map_err(to_pyvalue_err)
353 }
354
355 #[staticmethod]
356 #[pyo3(name = "from_str")]
357 fn py_from_str(value: &str) -> Self {
358 Self::from(value)
359 }
360
361 #[pyo3(name = "is_zero")]
362 fn py_is_zero(&self) -> bool {
363 self.is_zero()
364 }
365
366 #[pyo3(name = "is_positive")]
367 fn py_is_positive(&self) -> bool {
368 self.is_positive()
369 }
370
371 #[pyo3(name = "as_decimal")]
372 fn py_as_decimal(&self) -> Decimal {
373 self.as_decimal()
374 }
375
376 #[pyo3(name = "as_double")]
377 fn py_as_double(&self) -> f64 {
378 self.as_f64()
379 }
380
381 #[pyo3(name = "to_formatted_str")]
382 fn py_to_formatted_str(&self) -> String {
383 self.to_formatted_string()
384 }
385}