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