nautilus_model/python/events/account/
state.rs

1// -------------------------------------------------------------------------------------------------
2//  Copyright (C) 2015-2025 Nautech Systems Pty Ltd. All rights reserved.
3//  https://nautechsystems.io
4//
5//  Licensed under the GNU Lesser General Public License Version 3.0 (the "License");
6//  You may not use this file except in compliance with the License.
7//  You may obtain a copy of the License at https://www.gnu.org/licenses/lgpl-3.0.en.html
8//
9//  Unless required by applicable law or agreed to in writing, software
10//  distributed under the License is distributed on an "AS IS" BASIS,
11//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12//  See the License for the specific language governing permissions and
13//  limitations under the License.
14// -------------------------------------------------------------------------------------------------
15
16use 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        // iterate over balances and margins and run to_dict on each item and collect them
150        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}