nautilus_model/python/types/
quantity.rs

1// -------------------------------------------------------------------------------------------------
2//  Copyright (C) 2015-2025 Nautech Systems Pty Ltd. All rights reserved.
3//  https://nautechsystems.io
4//
5//  Licensed under the GNU Lesser General Public License Version 3.0 (the "License");
6//  You may not use this file except in compliance with the License.
7//  You may obtain a copy of the License at https://www.gnu.org/licenses/lgpl-3.0.en.html
8//
9//  Unless required by applicable law or agreed to in writing, software
10//  distributed under the License is distributed on an "AS IS" BASIS,
11//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12//  See the License for the specific language governing permissions and
13//  limitations under the License.
14// -------------------------------------------------------------------------------------------------
15
16use 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)) // Safe default
62    }
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}