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