nautilus_model/python/events/account/
state.rs1use std::str::FromStr;
17
18use nautilus_core::{
19 Params, UUID4,
20 python::{
21 IntoPyObjectNautilusExt,
22 params::params_to_pydict,
23 parsing::{get_required, get_required_list, get_required_parsed, get_required_string},
24 to_pyvalue_err,
25 },
26};
27use pyo3::{basic::CompareOp, prelude::*, types::PyDict};
28
29use crate::{
30 enums::AccountType,
31 events::AccountState,
32 identifiers::AccountId,
33 types::{AccountBalance, Currency, MarginBalance},
34};
35
36#[pyo3_stub_gen::derive::gen_stub_pymethods]
37#[pymethods]
38impl AccountState {
39 #[expect(clippy::too_many_arguments)]
46 #[new]
47 #[pyo3(signature = (account_id, account_type, balances, margins, is_reported, event_id, ts_event, ts_init, base_currency=None, info=None))]
48 fn py_new(
49 account_id: AccountId,
50 account_type: AccountType,
51 balances: Vec<AccountBalance>,
52 margins: Vec<MarginBalance>,
53 is_reported: bool,
54 event_id: UUID4,
55 ts_event: u64,
56 ts_init: u64,
57 base_currency: Option<Currency>,
58 info: Option<pyo3::Py<PyDict>>,
59 ) -> PyResult<Self> {
60 let info_params = info
61 .map(|dict| Python::attach(|py| nautilus_core::from_pydict(py, &dict)))
62 .transpose()?
63 .flatten();
64 Ok(Self::new(
65 account_id,
66 account_type,
67 balances,
68 margins,
69 is_reported,
70 event_id,
71 ts_event.into(),
72 ts_init.into(),
73 base_currency,
74 )
75 .with_info(info_params))
76 }
77
78 #[getter]
79 fn account_id(&self) -> AccountId {
80 self.account_id
81 }
82
83 #[getter]
84 fn account_type(&self) -> AccountType {
85 self.account_type
86 }
87
88 #[getter]
89 fn base_currency(&self) -> Option<Currency> {
90 self.base_currency
91 }
92
93 #[getter]
94 fn balances(&self) -> Vec<AccountBalance> {
95 self.balances.clone()
96 }
97
98 #[getter]
99 fn margins(&self) -> Vec<MarginBalance> {
100 self.margins.clone()
101 }
102
103 #[getter]
104 #[pyo3(name = "is_reported")]
105 fn py_is_reported(&self) -> bool {
106 self.is_reported
107 }
108
109 #[getter]
110 #[pyo3(name = "event_id")]
111 fn py_event_id(&self) -> UUID4 {
112 self.event_id
113 }
114
115 #[getter]
116 #[pyo3(name = "ts_event")]
117 fn py_ts_event(&self) -> u64 {
118 self.ts_event.as_u64()
119 }
120
121 #[getter]
122 #[pyo3(name = "ts_init")]
123 fn py_ts_init(&self) -> u64 {
124 self.ts_init.as_u64()
125 }
126
127 #[getter]
128 #[pyo3(name = "info")]
129 fn py_info(&self, py: Python<'_>) -> PyResult<Py<PyDict>> {
130 match &self.info {
131 Some(params) => params_to_pydict(py, params),
132 None => Ok(PyDict::new(py).unbind()),
133 }
134 }
135
136 fn __richcmp__(&self, other: &Self, op: CompareOp, py: Python<'_>) -> Py<PyAny> {
137 match op {
138 CompareOp::Eq => self.eq(other).into_py_any_unwrap(py),
139 CompareOp::Ne => self.ne(other).into_py_any_unwrap(py),
140 _ => py.NotImplemented(),
141 }
142 }
143
144 fn __repr__(&self) -> String {
145 format!("{self:?}")
146 }
147
148 fn __str__(&self) -> String {
149 self.to_string()
150 }
151
152 #[staticmethod]
153 #[pyo3(name = "from_dict")]
159 pub fn py_from_dict(values: &Bound<'_, PyDict>) -> PyResult<Self> {
160 let account_id = get_required_string(values, "account_id")?;
161 let _account_type = get_required_string(values, "account_type")?;
162 let base_currency_str = get_required_string(values, "base_currency")?;
163 let base_currency = if base_currency_str == "None" {
164 None
165 } else {
166 Some(
167 Currency::from_str(&base_currency_str)
168 .map_err(|e| to_pyvalue_err(format!("Failed to parse 'base_currency': {e}")))?,
169 )
170 };
171 let balances_list = get_required_list(values, "balances")?;
172 let balances: Vec<AccountBalance> = balances_list
173 .iter()
174 .map(|b| {
175 let balance_dict = b.extract::<Bound<'_, PyDict>>()?;
176 AccountBalance::py_from_dict(&balance_dict)
177 })
178 .collect::<PyResult<Vec<AccountBalance>>>()?;
179 let margins_list = get_required_list(values, "margins")?;
180 let margins: Vec<MarginBalance> = margins_list
181 .iter()
182 .map(|m| {
183 let margin_dict = m.extract::<Bound<'_, PyDict>>()?;
184 MarginBalance::py_from_dict(&margin_dict)
185 })
186 .collect::<PyResult<Vec<MarginBalance>>>()?;
187 let reported = get_required::<bool>(values, "reported")?;
188 let _event_id = get_required_string(values, "event_id")?;
189 let ts_event = get_required::<u64>(values, "ts_event")?;
190 let ts_init = get_required::<u64>(values, "ts_init")?;
191 let info: Option<Params> = match values.get_item("info")? {
192 Some(item) if !item.is_none() => {
193 let info_dict: Bound<'_, PyDict> = item.extract()?;
194 if info_dict.is_empty() {
195 None
196 } else {
197 nautilus_core::from_pydict(info_dict.py(), &info_dict.clone().unbind())?
198 }
199 }
200 _ => None,
201 };
202 let account = Self::new(
203 AccountId::from(account_id.as_str()),
204 get_required_parsed(values, "account_type", |s| {
205 AccountType::from_str(&s).map_err(|e| e.to_string())
206 })?,
207 balances,
208 margins,
209 reported,
210 get_required_parsed(values, "event_id", |s| UUID4::from_str(&s))?,
211 ts_event.into(),
212 ts_init.into(),
213 base_currency,
214 )
215 .with_info(info);
216 Ok(account)
217 }
218
219 #[pyo3(name = "to_dict")]
225 pub fn py_to_dict(&self, py: Python<'_>) -> PyResult<Py<PyAny>> {
226 let dict = PyDict::new(py);
227 dict.set_item("type", stringify!(AccountState))?;
228 dict.set_item("account_id", self.account_id.to_string())?;
229 dict.set_item("account_type", self.account_type.to_string())?;
230 let balances_dict: PyResult<Vec<_>> =
232 self.balances.iter().map(|b| b.py_to_dict(py)).collect();
233 let margins_dict: PyResult<Vec<_>> =
234 self.margins.iter().map(|m| m.py_to_dict(py)).collect();
235 dict.set_item("balances", balances_dict?)?;
236 dict.set_item("margins", margins_dict?)?;
237 dict.set_item("reported", self.is_reported)?;
238 dict.set_item("event_id", self.event_id.to_string())?;
239 match &self.info {
240 Some(params) => dict.set_item("info", params_to_pydict(py, params)?)?,
241 None => dict.set_item("info", PyDict::new(py))?,
242 }
243 dict.set_item("ts_event", self.ts_event.as_u64())?;
244 dict.set_item("ts_init", self.ts_init.as_u64())?;
245
246 match self.base_currency {
247 Some(base_currency) => {
248 dict.set_item("base_currency", base_currency.code.to_string())?;
249 }
250 None => dict.set_item("base_currency", "None")?,
251 }
252 Ok(dict.into())
253 }
254}