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