nautilus_model/python/events/account/
state.rs1use std::str::FromStr;
17
18use nautilus_core::{
19 UUID4,
20 python::{
21 IntoPyObjectNautilusExt,
22 parsing::{get_required, get_required_list, get_required_parsed, get_required_string},
23 },
24};
25use pyo3::{basic::CompareOp, prelude::*, types::PyDict};
26
27use crate::{
28 enums::AccountType,
29 events::AccountState,
30 identifiers::AccountId,
31 types::{AccountBalance, Currency, MarginBalance},
32};
33
34#[pymethods]
35#[pyo3_stub_gen::derive::gen_stub_pymethods]
36impl AccountState {
37 #[expect(clippy::too_many_arguments)]
39 #[new]
40 #[pyo3(signature = (account_id, account_type, balances, margins, is_reported, event_id, ts_event, ts_init, base_currency=None))]
41 fn py_new(
42 account_id: AccountId,
43 account_type: AccountType,
44 balances: Vec<AccountBalance>,
45 margins: Vec<MarginBalance>,
46 is_reported: bool,
47 event_id: UUID4,
48 ts_event: u64,
49 ts_init: u64,
50 base_currency: Option<Currency>,
51 ) -> Self {
52 Self::new(
53 account_id,
54 account_type,
55 balances,
56 margins,
57 is_reported,
58 event_id,
59 ts_event.into(),
60 ts_init.into(),
61 base_currency,
62 )
63 }
64
65 #[getter]
66 fn account_id(&self) -> AccountId {
67 self.account_id
68 }
69
70 #[getter]
71 fn account_type(&self) -> AccountType {
72 self.account_type
73 }
74
75 #[getter]
76 fn base_currency(&self) -> Option<Currency> {
77 self.base_currency
78 }
79
80 #[getter]
81 fn balances(&self) -> Vec<AccountBalance> {
82 self.balances.clone()
83 }
84
85 #[getter]
86 fn margins(&self) -> Vec<MarginBalance> {
87 self.margins.clone()
88 }
89
90 fn __richcmp__(&self, other: &Self, op: CompareOp, py: Python<'_>) -> Py<PyAny> {
91 match op {
92 CompareOp::Eq => self.eq(other).into_py_any_unwrap(py),
93 CompareOp::Ne => self.ne(other).into_py_any_unwrap(py),
94 _ => py.NotImplemented(),
95 }
96 }
97
98 fn __repr__(&self) -> String {
99 format!("{self:?}")
100 }
101
102 fn __str__(&self) -> String {
103 self.to_string()
104 }
105
106 #[staticmethod]
107 #[pyo3(name = "from_dict")]
114 pub fn py_from_dict(values: &Bound<'_, PyDict>) -> PyResult<Self> {
115 let account_id = get_required_string(values, "account_id")?;
116 let _account_type = get_required_string(values, "account_type")?;
117 let _base_currency = get_required_string(values, "base_currency")?;
118 let balances_list = get_required_list(values, "balances")?;
119 let balances: Vec<AccountBalance> = balances_list
120 .iter()
121 .map(|b| {
122 let balance_dict = b.extract::<Bound<'_, PyDict>>()?;
123 AccountBalance::py_from_dict(&balance_dict)
124 })
125 .collect::<PyResult<Vec<AccountBalance>>>()?;
126 let margins_list = get_required_list(values, "margins")?;
127 let margins: Vec<MarginBalance> = margins_list
128 .iter()
129 .map(|m| {
130 let margin_dict = m.extract::<Bound<'_, PyDict>>()?;
131 MarginBalance::py_from_dict(&margin_dict)
132 })
133 .collect::<PyResult<Vec<MarginBalance>>>()?;
134 let reported = get_required::<bool>(values, "reported")?;
135 let _event_id = get_required_string(values, "event_id")?;
136 let ts_event = get_required::<u64>(values, "ts_event")?;
137 let ts_init = get_required::<u64>(values, "ts_init")?;
138 let account = Self::new(
139 AccountId::from(account_id.as_str()),
140 get_required_parsed(values, "account_type", |s| {
141 AccountType::from_str(&s).map_err(|e| e.to_string())
142 })?,
143 balances,
144 margins,
145 reported,
146 get_required_parsed(values, "event_id", |s| UUID4::from_str(&s))?,
147 ts_event.into(),
148 ts_init.into(),
149 Some(get_required_parsed(values, "base_currency", |s| {
150 Currency::from_str(&s).map_err(|e| e.to_string())
151 })?),
152 );
153 Ok(account)
154 }
155
156 #[pyo3(name = "to_dict")]
162 pub fn py_to_dict(&self, py: Python<'_>) -> PyResult<Py<PyAny>> {
163 let dict = PyDict::new(py);
164 dict.set_item("type", stringify!(AccountState))?;
165 dict.set_item("account_id", self.account_id.to_string())?;
166 dict.set_item("account_type", self.account_type.to_string())?;
167 let balances_dict: PyResult<Vec<_>> =
169 self.balances.iter().map(|b| b.py_to_dict(py)).collect();
170 let margins_dict: PyResult<Vec<_>> =
171 self.margins.iter().map(|m| m.py_to_dict(py)).collect();
172 dict.set_item("balances", balances_dict?)?;
173 dict.set_item("margins", margins_dict?)?;
174 dict.set_item("reported", self.is_reported)?;
175 dict.set_item("event_id", self.event_id.to_string())?;
176 dict.set_item("info", PyDict::new(py))?;
177 dict.set_item("ts_event", self.ts_event.as_u64())?;
178 dict.set_item("ts_init", self.ts_init.as_u64())?;
179
180 match self.base_currency {
181 Some(base_currency) => {
182 dict.set_item("base_currency", base_currency.code.to_string())?;
183 }
184 None => dict.set_item("base_currency", "None")?,
185 }
186 Ok(dict.into())
187 }
188}