nautilus_indicators/python/average/
wma.rs1use nautilus_core::python::{to_pytype_err, to_pyvalue_err};
17use nautilus_model::{
18 data::{Bar, QuoteTick, TradeTick},
19 enums::PriceType,
20};
21use pyo3::{ffi, prelude::*, types::PyString};
22
23use crate::{
24 average::wma::WeightedMovingAverage,
25 indicator::{Indicator, MovingAverage},
26 python::float_precision,
27};
28
29#[pyo3_stub_gen::derive::gen_stub_pymethods]
30#[pymethods]
31impl WeightedMovingAverage {
32 #[new]
34 #[pyo3(signature = (period, weights=None, price_type=None))]
35 pub fn py_new(
36 period: usize,
37 #[gen_stub(override_type(type_repr = "typing.Sequence[float] | None"))] weights: Option<
38 &Bound<'_, PyAny>,
39 >,
40 price_type: Option<PriceType>,
41 ) -> PyResult<Self> {
42 Self::check_period(period).map_err(to_pyvalue_err)?;
43
44 match weights {
45 Some(weights) => {
46 if weights.is_instance_of::<PyString>() {
47 return Err(to_pytype_err("Can't extract `str` to `Vec`"));
48 }
49
50 #[allow(
51 unsafe_code,
52 reason = "preserve PyO3 Vec extraction's sequence predicate for custom Python sequences"
53 )]
54 let is_sequence = unsafe { ffi::PySequence_Check(weights.as_ptr()) } != 0;
56 if !is_sequence {
57 return Err(to_pytype_err("weights must be a sequence"));
58 }
59
60 let mut values = Vec::new();
61
62 for item in weights.try_iter()?.take(period + 1) {
64 values.try_reserve(1).map_err(to_pyvalue_err)?;
65 values.push(item?.extract::<f64>()?);
66 }
67
68 Self::with_weights_checked(period, values, price_type)
69 }
70 None => Self::new_checked(period, price_type),
71 }
72 .map_err(to_pyvalue_err)
73 }
74
75 fn __repr__(&self) -> String {
76 format!("WeightedMovingAverage({},{:?})", self.period, self.weights)
77 }
78
79 #[getter]
80 #[pyo3(name = "name")]
81 fn py_name(&self) -> String {
82 self.name()
83 }
84
85 #[getter]
86 #[pyo3(name = "period")]
87 const fn py_period(&self) -> usize {
88 self.period
89 }
90
91 #[getter]
92 #[pyo3(name = "price_type")]
93 const fn py_price_type(&self) -> PriceType {
94 self.price_type
95 }
96
97 #[getter]
98 #[pyo3(name = "value")]
99 const fn py_value(&self) -> f64 {
100 self.value
101 }
102
103 #[getter]
104 #[pyo3(name = "weights")]
105 fn py_weights(&self) -> Vec<f64> {
106 self.weights.clone()
107 }
108
109 #[getter]
110 #[pyo3(name = "count")]
111 fn py_count(&self) -> usize {
112 self.count()
113 }
114
115 #[getter]
116 #[pyo3(name = "has_inputs")]
117 fn py_has_inputs(&self) -> bool {
118 self.has_inputs()
119 }
120
121 #[getter]
122 #[pyo3(name = "initialized")]
123 const fn py_initialized(&self) -> bool {
124 self.initialized
125 }
126
127 #[pyo3(name = "handle_quote_tick")]
128 fn py_handle_quote_tick(&mut self, quote: &QuoteTick) -> PyResult<()> {
129 float_precision::check_quote(quote)?;
130 self.handle_quote(quote).map_err(to_pyvalue_err)
131 }
132
133 #[pyo3(name = "handle_trade_tick")]
134 fn py_handle_trade_tick(&mut self, trade: &TradeTick) -> PyResult<()> {
135 float_precision::check_trade(trade)?;
136 self.handle_trade(trade);
137 Ok(())
138 }
139
140 #[pyo3(name = "handle_bar")]
141 fn py_handle_bar(&mut self, bar: &Bar) -> PyResult<()> {
142 float_precision::check_bar(bar)?;
143 self.handle_bar(bar);
144 Ok(())
145 }
146
147 #[pyo3(name = "reset")]
148 fn py_reset(&mut self) {
149 self.reset();
150 }
151
152 #[pyo3(name = "update_raw")]
153 fn py_update_raw(&mut self, value: f64) {
154 self.update_raw(value);
155 }
156}