1use std::{
17 collections::hash_map::DefaultHasher,
18 hash::{Hash, Hasher},
19 ops::Neg,
20 str::FromStr,
21};
22
23use nautilus_core::python::{
24 correctness_error_to_pyvalue_err, get_pytype_name, to_pytype_err, to_pyvalue_err,
25};
26use pyo3::{basic::CompareOp, conversion::IntoPyObjectExt, prelude::*, types::PyFloat};
27use rust_decimal::{Decimal, RoundingStrategy};
28
29use super::fixed::{
30 ArithmeticError, ArithmeticOperation, FloatArithmetic, check_raw_scales,
31 extract_arithmetic_decimal,
32};
33use crate::types::{Quantity, fixed::raw_scale, quantity::QuantityRaw};
34
35#[pymethods]
36#[pyo3_stub_gen::derive::gen_stub_pymethods]
37impl Quantity {
38 #[new]
49 fn py_new(value: f64, precision: u8) -> PyResult<Self> {
50 Self::new_checked(value, precision).map_err(to_pyvalue_err)
51 }
52
53 fn __reduce__(&self, py: Python) -> PyResult<Py<PyAny>> {
54 let from_raw = py.get_type::<Self>().getattr("from_raw")?;
55 let args = (self.raw(), self.precision).into_py_any(py)?;
56 (from_raw, args).into_py_any(py)
57 }
58
59 fn __richcmp__(
60 &self,
61 other: &Bound<'_, PyAny>,
62 op: CompareOp,
63 py: Python<'_>,
64 ) -> PyResult<Py<PyAny>> {
65 if let Ok(other_qty) = other.extract::<Self>() {
66 let result = match op {
67 CompareOp::Eq => self.eq(&other_qty),
68 CompareOp::Ne => self.ne(&other_qty),
69 CompareOp::Ge => self.ge(&other_qty),
70 CompareOp::Gt => self.gt(&other_qty),
71 CompareOp::Le => self.le(&other_qty),
72 CompareOp::Lt => self.lt(&other_qty),
73 };
74 result.into_py_any(py)
75 } else if let Ok(other_dec) = other.extract::<Decimal>() {
76 let result = match op {
77 CompareOp::Eq => self.as_decimal() == other_dec,
78 CompareOp::Ne => self.as_decimal() != other_dec,
79 CompareOp::Ge => self.as_decimal() >= other_dec,
80 CompareOp::Gt => self.as_decimal() > other_dec,
81 CompareOp::Le => self.as_decimal() <= other_dec,
82 CompareOp::Lt => self.as_decimal() < other_dec,
83 };
84 result.into_py_any(py)
85 } else {
86 Ok(py.NotImplemented())
87 }
88 }
89
90 fn __hash__(&self) -> isize {
91 let mut h = DefaultHasher::new();
92 self.hash(&mut h);
93 h.finish() as isize
94 }
95
96 fn __add__(&self, other: &Bound<'_, PyAny>, py: Python) -> PyResult<Py<PyAny>> {
97 if other.is_instance_of::<PyFloat>() {
98 let other_float: f64 = other.extract()?;
99 ArithmeticOperation::Add
100 .checked_f64(self.as_f64_checked()?, other_float)?
101 .into_py_any(py)
102 } else if let Ok(other_qty) = other.extract::<Self>() {
103 check_raw_scales(self.precision, other_qty.precision)?;
104 (*self)
105 .checked_add(other_qty)
106 .ok_or(ArithmeticError::Overflow)?
107 .into_py_any(py)
108 } else if let Some(other_dec) = extract_arithmetic_decimal(other) {
109 ArithmeticOperation::Add
110 .checked_decimal(self.as_decimal(), other_dec)?
111 .into_py_any(py)
112 } else {
113 let pytype_name = get_pytype_name(other)?;
114 Err(to_pytype_err(format!(
115 "Unsupported type for __add__, was `{pytype_name}`"
116 )))
117 }
118 }
119
120 fn __radd__(&self, other: &Bound<'_, PyAny>, py: Python) -> PyResult<Py<PyAny>> {
121 if other.is_instance_of::<PyFloat>() {
122 let other_float: f64 = other.extract()?;
123 ArithmeticOperation::Add
124 .checked_f64(other_float, self.as_f64_checked()?)?
125 .into_py_any(py)
126 } else if let Ok(other_qty) = other.extract::<Self>() {
127 check_raw_scales(other_qty.precision, self.precision)?;
128 other_qty
129 .checked_add(*self)
130 .ok_or(ArithmeticError::Overflow)?
131 .into_py_any(py)
132 } else if let Some(other_dec) = extract_arithmetic_decimal(other) {
133 ArithmeticOperation::Add
134 .checked_decimal(other_dec, self.as_decimal())?
135 .into_py_any(py)
136 } else {
137 let pytype_name = get_pytype_name(other)?;
138 Err(to_pytype_err(format!(
139 "Unsupported type for __radd__, was `{pytype_name}`"
140 )))
141 }
142 }
143
144 fn __sub__(&self, other: &Bound<'_, PyAny>, py: Python) -> PyResult<Py<PyAny>> {
145 if other.is_instance_of::<PyFloat>() {
146 let other_float: f64 = other.extract()?;
147 ArithmeticOperation::Sub
148 .checked_f64(self.as_f64_checked()?, other_float)?
149 .into_py_any(py)
150 } else if let Ok(other_qty) = other.extract::<Self>() {
151 check_raw_scales(self.precision, other_qty.precision)?;
152 if other_qty > *self {
153 return Err(to_pyvalue_err(format!(
154 "Quantity subtraction would result in negative value: {self} - {other_qty}"
155 )));
156 }
157 (*self)
158 .checked_sub(other_qty)
159 .ok_or(ArithmeticError::Overflow)?
160 .into_py_any(py)
161 } else if let Some(other_dec) = extract_arithmetic_decimal(other) {
162 ArithmeticOperation::Sub
163 .checked_decimal(self.as_decimal(), other_dec)?
164 .into_py_any(py)
165 } else {
166 let pytype_name = get_pytype_name(other)?;
167 Err(to_pytype_err(format!(
168 "Unsupported type for __sub__, was `{pytype_name}`"
169 )))
170 }
171 }
172
173 fn __rsub__(&self, other: &Bound<'_, PyAny>, py: Python) -> PyResult<Py<PyAny>> {
174 if other.is_instance_of::<PyFloat>() {
175 let other_float: f64 = other.extract()?;
176 ArithmeticOperation::Sub
177 .checked_f64(other_float, self.as_f64_checked()?)?
178 .into_py_any(py)
179 } else if let Ok(other_qty) = other.extract::<Self>() {
180 check_raw_scales(other_qty.precision, self.precision)?;
181 if *self > other_qty {
182 return Err(to_pyvalue_err(format!(
183 "Quantity subtraction would result in negative value: {other_qty} - {self}"
184 )));
185 }
186 other_qty
187 .checked_sub(*self)
188 .ok_or(ArithmeticError::Overflow)?
189 .into_py_any(py)
190 } else if let Some(other_dec) = extract_arithmetic_decimal(other) {
191 ArithmeticOperation::Sub
192 .checked_decimal(other_dec, self.as_decimal())?
193 .into_py_any(py)
194 } else {
195 let pytype_name = get_pytype_name(other)?;
196 Err(to_pytype_err(format!(
197 "Unsupported type for __rsub__, was `{pytype_name}`"
198 )))
199 }
200 }
201
202 fn __mul__(&self, other: &Bound<'_, PyAny>, py: Python) -> PyResult<Py<PyAny>> {
203 if other.is_instance_of::<PyFloat>() {
204 let other_float: f64 = other.extract()?;
205 ArithmeticOperation::Mul
206 .checked_f64(self.as_f64_checked()?, other_float)?
207 .into_py_any(py)
208 } else if let Ok(other_qty) = other.extract::<Self>() {
209 ArithmeticOperation::Mul
210 .checked_decimal(self.as_decimal(), other_qty.as_decimal())?
211 .into_py_any(py)
212 } else if let Some(other_dec) = extract_arithmetic_decimal(other) {
213 ArithmeticOperation::Mul
214 .checked_decimal(self.as_decimal(), other_dec)?
215 .into_py_any(py)
216 } else {
217 let pytype_name = get_pytype_name(other)?;
218 Err(to_pytype_err(format!(
219 "Unsupported type for __mul__, was `{pytype_name}`"
220 )))
221 }
222 }
223
224 fn __rmul__(&self, other: &Bound<'_, PyAny>, py: Python) -> PyResult<Py<PyAny>> {
225 if other.is_instance_of::<PyFloat>() {
226 let other_float: f64 = other.extract()?;
227 ArithmeticOperation::Mul
228 .checked_f64(other_float, self.as_f64_checked()?)?
229 .into_py_any(py)
230 } else if let Ok(other_qty) = other.extract::<Self>() {
231 ArithmeticOperation::Mul
232 .checked_decimal(other_qty.as_decimal(), self.as_decimal())?
233 .into_py_any(py)
234 } else if let Some(other_dec) = extract_arithmetic_decimal(other) {
235 ArithmeticOperation::Mul
236 .checked_decimal(other_dec, self.as_decimal())?
237 .into_py_any(py)
238 } else {
239 let pytype_name = get_pytype_name(other)?;
240 Err(to_pytype_err(format!(
241 "Unsupported type for __rmul__, was `{pytype_name}`"
242 )))
243 }
244 }
245
246 fn __truediv__(&self, other: &Bound<'_, PyAny>, py: Python) -> PyResult<Py<PyAny>> {
247 if other.is_instance_of::<PyFloat>() {
248 let other_float: f64 = other.extract()?;
249 ArithmeticOperation::Div
250 .checked_f64(self.as_f64_checked()?, other_float)?
251 .into_py_any(py)
252 } else if let Ok(other_qty) = other.extract::<Self>() {
253 ArithmeticOperation::Div
254 .checked_decimal(self.as_decimal(), other_qty.as_decimal())?
255 .into_py_any(py)
256 } else if let Some(other_dec) = extract_arithmetic_decimal(other) {
257 ArithmeticOperation::Div
258 .checked_decimal(self.as_decimal(), other_dec)?
259 .into_py_any(py)
260 } else {
261 let pytype_name = get_pytype_name(other)?;
262 Err(to_pytype_err(format!(
263 "Unsupported type for __truediv__, was `{pytype_name}`"
264 )))
265 }
266 }
267
268 fn __rtruediv__(&self, other: &Bound<'_, PyAny>, py: Python) -> PyResult<Py<PyAny>> {
269 if other.is_instance_of::<PyFloat>() {
270 let other_float: f64 = other.extract()?;
271 ArithmeticOperation::Div
272 .checked_f64(other_float, self.as_f64_checked()?)?
273 .into_py_any(py)
274 } else if let Ok(other_qty) = other.extract::<Self>() {
275 ArithmeticOperation::Div
276 .checked_decimal(other_qty.as_decimal(), self.as_decimal())?
277 .into_py_any(py)
278 } else if let Some(other_dec) = extract_arithmetic_decimal(other) {
279 ArithmeticOperation::Div
280 .checked_decimal(other_dec, self.as_decimal())?
281 .into_py_any(py)
282 } else {
283 let pytype_name = get_pytype_name(other)?;
284 Err(to_pytype_err(format!(
285 "Unsupported type for __rtruediv__, was `{pytype_name}`"
286 )))
287 }
288 }
289
290 fn __floordiv__(&self, other: &Bound<'_, PyAny>, py: Python) -> PyResult<Py<PyAny>> {
291 if other.is_instance_of::<PyFloat>() {
292 let other_float: f64 = other.extract()?;
293 ArithmeticOperation::Div
294 .checked_f64(self.as_f64_checked()?, other_float)?
295 .floor()
296 .into_py_any(py)
297 } else if let Ok(other_qty) = other.extract::<Self>() {
298 ArithmeticOperation::Div
299 .checked_decimal(self.as_decimal(), other_qty.as_decimal())?
300 .floor()
301 .into_py_any(py)
302 } else if let Some(other_dec) = extract_arithmetic_decimal(other) {
303 ArithmeticOperation::Div
304 .checked_decimal(self.as_decimal(), other_dec)?
305 .floor()
306 .into_py_any(py)
307 } else {
308 let pytype_name = get_pytype_name(other)?;
309 Err(to_pytype_err(format!(
310 "Unsupported type for __floordiv__, was `{pytype_name}`"
311 )))
312 }
313 }
314
315 fn __rfloordiv__(&self, other: &Bound<'_, PyAny>, py: Python) -> PyResult<Py<PyAny>> {
316 if other.is_instance_of::<PyFloat>() {
317 let other_float: f64 = other.extract()?;
318 ArithmeticOperation::Div
319 .checked_f64(other_float, self.as_f64_checked()?)?
320 .floor()
321 .into_py_any(py)
322 } else if let Ok(other_qty) = other.extract::<Self>() {
323 ArithmeticOperation::Div
324 .checked_decimal(other_qty.as_decimal(), self.as_decimal())?
325 .floor()
326 .into_py_any(py)
327 } else if let Some(other_dec) = extract_arithmetic_decimal(other) {
328 ArithmeticOperation::Div
329 .checked_decimal(other_dec, self.as_decimal())?
330 .floor()
331 .into_py_any(py)
332 } else {
333 let pytype_name = get_pytype_name(other)?;
334 Err(to_pytype_err(format!(
335 "Unsupported type for __rfloordiv__, was `{pytype_name}`"
336 )))
337 }
338 }
339
340 fn __mod__(&self, other: &Bound<'_, PyAny>, py: Python) -> PyResult<Py<PyAny>> {
341 if other.is_instance_of::<PyFloat>() {
342 let other_float: f64 = other.extract()?;
343 ArithmeticOperation::Rem
344 .checked_f64(self.as_f64_checked()?, other_float)?
345 .into_py_any(py)
346 } else if let Ok(other_qty) = other.extract::<Self>() {
347 ArithmeticOperation::Rem
348 .checked_decimal(self.as_decimal(), other_qty.as_decimal())?
349 .into_py_any(py)
350 } else if let Some(other_dec) = extract_arithmetic_decimal(other) {
351 ArithmeticOperation::Rem
352 .checked_decimal(self.as_decimal(), other_dec)?
353 .into_py_any(py)
354 } else {
355 let pytype_name = get_pytype_name(other)?;
356 Err(to_pytype_err(format!(
357 "Unsupported type for __mod__, was `{pytype_name}`"
358 )))
359 }
360 }
361
362 fn __rmod__(&self, other: &Bound<'_, PyAny>, py: Python) -> PyResult<Py<PyAny>> {
363 if other.is_instance_of::<PyFloat>() {
364 let other_float: f64 = other.extract()?;
365 ArithmeticOperation::Rem
366 .checked_f64(other_float, self.as_f64_checked()?)?
367 .into_py_any(py)
368 } else if let Ok(other_qty) = other.extract::<Self>() {
369 ArithmeticOperation::Rem
370 .checked_decimal(other_qty.as_decimal(), self.as_decimal())?
371 .into_py_any(py)
372 } else if let Some(other_dec) = extract_arithmetic_decimal(other) {
373 ArithmeticOperation::Rem
374 .checked_decimal(other_dec, self.as_decimal())?
375 .into_py_any(py)
376 } else {
377 let pytype_name = get_pytype_name(other)?;
378 Err(to_pytype_err(format!(
379 "Unsupported type for __rmod__, was `{pytype_name}`"
380 )))
381 }
382 }
383
384 fn __neg__(&self) -> Decimal {
385 self.as_decimal().neg()
386 }
387
388 fn __pos__(&self) -> Self {
389 *self
390 }
391
392 fn __abs__(&self) -> Self {
393 *self
394 }
395
396 fn __int__(&self) -> QuantityRaw {
397 let scale = QuantityRaw::try_from(raw_scale(self.precision))
398 .expect("effective raw scale should fit in QuantityRaw");
399 self.raw() / scale
400 }
401
402 fn __float__(&self) -> PyResult<f64> {
403 self.as_f64_checked()
404 }
405
406 #[pyo3(signature = (ndigits=None))]
407 fn __round__(&self, ndigits: Option<u32>) -> Decimal {
408 self.as_decimal()
409 .round_dp_with_strategy(ndigits.unwrap_or(0), RoundingStrategy::MidpointNearestEven)
410 }
411
412 fn __repr__(&self) -> String {
413 format!("{self:?}")
414 }
415
416 fn __str__(&self) -> String {
417 self.to_string()
418 }
419
420 #[getter(raw)]
433 fn py_raw(&self) -> QuantityRaw {
434 self.raw()
435 }
436
437 #[getter]
438 fn precision(&self) -> u8 {
439 self.precision
440 }
441
442 #[staticmethod]
444 #[pyo3(name = "from_raw")]
445 fn py_from_raw(raw: QuantityRaw, precision: u8) -> PyResult<Self> {
446 Self::from_raw_checked(raw, precision).map_err(correctness_error_to_pyvalue_err)
447 }
448
449 #[staticmethod]
451 #[pyo3(name = "zero")]
452 #[pyo3(signature = (precision = 0))]
453 fn py_zero(precision: u8) -> PyResult<Self> {
454 Self::new_checked(0.0, precision).map_err(to_pyvalue_err)
455 }
456
457 #[staticmethod]
458 #[pyo3(name = "from_int")]
459 fn py_from_int(value: u64) -> PyResult<Self> {
460 Self::new_checked(value as f64, 0).map_err(to_pyvalue_err)
461 }
462
463 #[staticmethod]
464 #[pyo3(name = "from_str")]
465 fn py_from_str(value: &str) -> PyResult<Self> {
466 Self::from_str(value).map_err(to_pyvalue_err)
467 }
468
469 #[staticmethod]
481 #[pyo3(name = "from_decimal")]
482 fn py_from_decimal(decimal: Decimal) -> PyResult<Self> {
483 Self::from_decimal(decimal).map_err(to_pyvalue_err)
484 }
485
486 #[staticmethod]
499 #[pyo3(name = "from_decimal_dp")]
500 fn py_from_decimal_dp(decimal: Decimal, precision: u8) -> PyResult<Self> {
501 Self::from_decimal_dp(decimal, precision).map_err(to_pyvalue_err)
502 }
503
504 #[staticmethod]
509 #[pyo3(name = "from_mantissa_exponent")]
510 fn py_from_mantissa_exponent(mantissa: u64, exponent: i8, precision: u8) -> PyResult<Self> {
511 Self::from_mantissa_exponent_checked(mantissa, exponent, precision)
512 .map_err(correctness_error_to_pyvalue_err)
513 }
514
515 #[pyo3(name = "is_zero")]
517 fn py_is_zero(&self) -> bool {
518 self.is_zero()
519 }
520
521 #[pyo3(name = "is_positive")]
523 fn py_is_positive(&self) -> bool {
524 self.is_positive()
525 }
526
527 #[pyo3(name = "as_decimal")]
529 fn py_as_decimal(&self) -> Decimal {
530 self.as_decimal()
531 }
532
533 #[pyo3(name = "as_double")]
534 fn py_as_double(&self) -> PyResult<f64> {
535 self.as_f64_checked()
536 }
537
538 #[pyo3(name = "to_formatted_str")]
539 fn py_to_formatted_str(&self) -> String {
540 self.to_formatted_string()
541 }
542
543 #[pyo3(name = "saturating_sub")]
551 fn py_saturating_sub(&self, other: Self) -> PyResult<Self> {
552 check_raw_scales(self.precision, other.precision)?;
553 Ok(self.saturating_sub(other))
554 }
555
556 #[pyo3(name = "checked_add")]
563 fn py_checked_add(&self, other: Self) -> Option<Self> {
564 self.checked_add(other)
565 }
566
567 #[pyo3(name = "checked_sub")]
573 fn py_checked_sub(&self, other: Self) -> Option<Self> {
574 self.checked_sub(other)
575 }
576}
577
578#[cfg(test)]
579mod tests {
580 use pyo3::Python;
581 use rstest::rstest;
582
583 use super::*;
584 use crate::types::{fixed::FIXED_PRECISION, quantity::QUANTITY_RAW_MAX};
585
586 #[rstest]
587 #[cfg(feature = "defi")]
588 fn test_saturating_sub_rejects_mixed_scales_without_panicking() {
589 Python::initialize();
590 Python::attach(|py| {
591 let lhs = Quantity::from_raw(10_u128.pow(18), 18);
592 let rhs = Quantity::from_raw(10_u128.pow(16), 16);
593 let error = lhs.py_saturating_sub(rhs).unwrap_err();
594 assert!(error.is_instance_of::<pyo3::exceptions::PyValueError>(py));
595 });
596 }
597
598 #[rstest]
599 #[case("0", 0)]
600 #[case("0.000000001", 0)]
601 #[case("1.999999999", 1)]
602 #[case("50.25", 50)]
603 #[case("9007199253.999999999", 9_007_199_253)]
604 fn test_int_uses_exact_raw_value(#[case] value: &str, #[case] expected: u64) {
605 let quantity = Quantity::from_str(value).unwrap();
606
607 assert_eq!(quantity.__int__(), QuantityRaw::from(expected));
608 }
609
610 #[rstest]
611 fn test_int_preserves_domain_maximum() {
612 let expected: QuantityRaw = if cfg!(feature = "high-precision") {
613 34_028_236_692_093
614 } else {
615 18_446_744_073
616 };
617
618 assert_eq!(
619 Quantity::from_raw(QUANTITY_RAW_MAX, FIXED_PRECISION).__int__(),
620 expected
621 );
622 }
623
624 #[rstest]
625 #[cfg(feature = "defi")]
626 fn test_int_uses_defi_raw_scale() {
627 let quantity = Quantity::from_raw(1_999_999_999_999_999_999, crate::defi::WEI_PRECISION);
628
629 assert_eq!(quantity.__int__(), 1);
630 }
631
632 #[rstest]
633 fn test_py_from_raw_rejects_out_of_range_raw_value() {
634 Python::initialize();
635 Python::attach(|_| {
636 let raw = QUANTITY_RAW_MAX.saturating_add(1);
637 let error = Quantity::py_from_raw(raw, 0).unwrap_err();
638
639 assert_eq!(
640 error.to_string(),
641 format!("ValueError: raw value {raw} exceeds QUANTITY_RAW_MAX={QUANTITY_RAW_MAX}")
642 );
643 });
644 }
645
646 #[rstest]
647 fn test_py_from_mantissa_exponent_handles_precision_and_overflow() {
648 Python::initialize();
649 Python::attach(|_| {
650 #[cfg(feature = "defi")]
651 let max_precision = crate::defi::WEI_PRECISION;
652 #[cfg(not(feature = "defi"))]
653 let max_precision = FIXED_PRECISION;
654
655 let exponent = -i8::try_from(max_precision).unwrap();
656 let quantity = Quantity::py_from_mantissa_exponent(1, exponent, max_precision).unwrap();
657 let invalid_precision = max_precision + 1;
658 let precision_error =
659 Quantity::py_from_mantissa_exponent(1, 0, invalid_precision).unwrap_err();
660 let overflow_error = Quantity::py_from_mantissa_exponent(u64::MAX, 100, 0).unwrap_err();
661 let precision_name = if cfg!(feature = "defi") {
662 "WEI_PRECISION"
663 } else {
664 "FIXED_PRECISION"
665 };
666
667 assert_eq!(quantity.raw(), 1);
668 assert_eq!(quantity.precision, max_precision);
669 assert_eq!(
670 precision_error.to_string(),
671 format!(
672 "ValueError: `precision` exceeded maximum `{precision_name}` ({max_precision}), was {invalid_precision}"
673 )
674 );
675 assert_eq!(
676 overflow_error.to_string(),
677 "ValueError: Overflow in Quantity::from_mantissa_exponent \
678 (mantissa=18446744073709551615, exponent=100, precision=0)"
679 );
680 });
681 }
682}