Skip to main content

nautilus_model/python/events/account/
state.rs

1// -------------------------------------------------------------------------------------------------
2//  Copyright (C) 2015-2026 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::{
19    Params, UUID4,
20    python::{
21        IntoPyObjectNautilusExt,
22        params::params_to_pydict,
23        parsing::{get_required, get_required_list, get_required_parsed, get_required_string},
24        to_pyvalue_err,
25    },
26};
27use pyo3::{basic::CompareOp, prelude::*, types::PyDict};
28
29use crate::{
30    enums::AccountType,
31    events::AccountState,
32    identifiers::AccountId,
33    types::{AccountBalance, Currency, MarginBalance},
34};
35
36#[pyo3_stub_gen::derive::gen_stub_pymethods]
37#[pymethods]
38impl AccountState {
39    /// Represents an event which includes information on the state of the account.
40    ///
41    /// The optional `info` bag carries venue-specific account data that does not map
42    /// to the typed `balances` and `margins` fields, such as wallet balance,
43    /// available balance, or an account summary, so consumers can read the venue
44    /// context that accompanied a given snapshot.
45    #[expect(clippy::too_many_arguments)]
46    #[new]
47    #[pyo3(signature = (account_id, account_type, balances, margins, is_reported, event_id, ts_event, ts_init, base_currency=None, info=None))]
48    fn py_new(
49        account_id: AccountId,
50        account_type: AccountType,
51        balances: Vec<AccountBalance>,
52        margins: Vec<MarginBalance>,
53        is_reported: bool,
54        event_id: UUID4,
55        ts_event: u64,
56        ts_init: u64,
57        base_currency: Option<Currency>,
58        info: Option<pyo3::Py<PyDict>>,
59    ) -> PyResult<Self> {
60        let info_params = info
61            .map(|dict| Python::attach(|py| nautilus_core::from_pydict(py, &dict)))
62            .transpose()?
63            .flatten();
64        Ok(Self::new(
65            account_id,
66            account_type,
67            balances,
68            margins,
69            is_reported,
70            event_id,
71            ts_event.into(),
72            ts_init.into(),
73            base_currency,
74        )
75        .with_info(info_params))
76    }
77
78    #[getter]
79    fn account_id(&self) -> AccountId {
80        self.account_id
81    }
82
83    #[getter]
84    fn account_type(&self) -> AccountType {
85        self.account_type
86    }
87
88    #[getter]
89    fn base_currency(&self) -> Option<Currency> {
90        self.base_currency
91    }
92
93    #[getter]
94    fn balances(&self) -> Vec<AccountBalance> {
95        self.balances.clone()
96    }
97
98    #[getter]
99    fn margins(&self) -> Vec<MarginBalance> {
100        self.margins.clone()
101    }
102
103    #[getter]
104    #[pyo3(name = "is_reported")]
105    fn py_is_reported(&self) -> bool {
106        self.is_reported
107    }
108
109    #[getter]
110    #[pyo3(name = "event_id")]
111    fn py_event_id(&self) -> UUID4 {
112        self.event_id
113    }
114
115    #[getter]
116    #[pyo3(name = "ts_event")]
117    fn py_ts_event(&self) -> u64 {
118        self.ts_event.as_u64()
119    }
120
121    #[getter]
122    #[pyo3(name = "ts_init")]
123    fn py_ts_init(&self) -> u64 {
124        self.ts_init.as_u64()
125    }
126
127    #[getter]
128    #[pyo3(name = "info")]
129    fn py_info(&self, py: Python<'_>) -> PyResult<Py<PyDict>> {
130        match &self.info {
131            Some(params) => params_to_pydict(py, params),
132            None => Ok(PyDict::new(py).unbind()),
133        }
134    }
135
136    fn __richcmp__(&self, other: &Self, op: CompareOp, py: Python<'_>) -> Py<PyAny> {
137        match op {
138            CompareOp::Eq => self.eq(other).into_py_any_unwrap(py),
139            CompareOp::Ne => self.ne(other).into_py_any_unwrap(py),
140            _ => py.NotImplemented(),
141        }
142    }
143
144    fn __repr__(&self) -> String {
145        format!("{self:?}")
146    }
147
148    fn __str__(&self) -> String {
149        self.to_string()
150    }
151
152    #[staticmethod]
153    /// Constructs an [`AccountState`] from a Python dict.
154    ///
155    /// # Errors
156    ///
157    /// Returns a `PyErr` if any required field is missing or type conversion fails.
158    #[pyo3(name = "from_dict")]
159    pub fn py_from_dict(values: &Bound<'_, PyDict>) -> PyResult<Self> {
160        let account_id = get_required_string(values, "account_id")?;
161        let _account_type = get_required_string(values, "account_type")?;
162        let base_currency_str = get_required_string(values, "base_currency")?;
163        let base_currency = if base_currency_str == "None" {
164            None
165        } else {
166            Some(
167                Currency::from_str(&base_currency_str)
168                    .map_err(|e| to_pyvalue_err(format!("Failed to parse 'base_currency': {e}")))?,
169            )
170        };
171        let balances_list = get_required_list(values, "balances")?;
172        let balances: Vec<AccountBalance> = balances_list
173            .iter()
174            .map(|b| {
175                let balance_dict = b.extract::<Bound<'_, PyDict>>()?;
176                AccountBalance::py_from_dict(&balance_dict)
177            })
178            .collect::<PyResult<Vec<AccountBalance>>>()?;
179        let margins_list = get_required_list(values, "margins")?;
180        let margins: Vec<MarginBalance> = margins_list
181            .iter()
182            .map(|m| {
183                let margin_dict = m.extract::<Bound<'_, PyDict>>()?;
184                MarginBalance::py_from_dict(&margin_dict)
185            })
186            .collect::<PyResult<Vec<MarginBalance>>>()?;
187        let reported = get_required::<bool>(values, "reported")?;
188        let _event_id = get_required_string(values, "event_id")?;
189        let ts_event = get_required::<u64>(values, "ts_event")?;
190        let ts_init = get_required::<u64>(values, "ts_init")?;
191        let info: Option<Params> = match values.get_item("info")? {
192            Some(item) if !item.is_none() => {
193                let info_dict: Bound<'_, PyDict> = item.extract()?;
194                if info_dict.is_empty() {
195                    None
196                } else {
197                    nautilus_core::from_pydict(info_dict.py(), &info_dict.clone().unbind())?
198                }
199            }
200            _ => None,
201        };
202        let account = Self::new(
203            AccountId::from(account_id.as_str()),
204            get_required_parsed(values, "account_type", |s| {
205                AccountType::from_str(&s).map_err(|e| e.to_string())
206            })?,
207            balances,
208            margins,
209            reported,
210            get_required_parsed(values, "event_id", |s| UUID4::from_str(&s))?,
211            ts_event.into(),
212            ts_init.into(),
213            base_currency,
214        )
215        .with_info(info);
216        Ok(account)
217    }
218
219    /// Converts this [`AccountState`] into a Python dict.
220    ///
221    /// # Errors
222    ///
223    /// Returns a `PyErr` if serialization into a Python dict fails.
224    #[pyo3(name = "to_dict")]
225    pub fn py_to_dict(&self, py: Python<'_>) -> PyResult<Py<PyAny>> {
226        let dict = PyDict::new(py);
227        dict.set_item("type", stringify!(AccountState))?;
228        dict.set_item("account_id", self.account_id.to_string())?;
229        dict.set_item("account_type", self.account_type.to_string())?;
230        // iterate over balances and margins and run to_dict on each item and collect them
231        let balances_dict: PyResult<Vec<_>> =
232            self.balances.iter().map(|b| b.py_to_dict(py)).collect();
233        let margins_dict: PyResult<Vec<_>> =
234            self.margins.iter().map(|m| m.py_to_dict(py)).collect();
235        dict.set_item("balances", balances_dict?)?;
236        dict.set_item("margins", margins_dict?)?;
237        dict.set_item("reported", self.is_reported)?;
238        dict.set_item("event_id", self.event_id.to_string())?;
239        match &self.info {
240            Some(params) => dict.set_item("info", params_to_pydict(py, params)?)?,
241            None => dict.set_item("info", PyDict::new(py))?,
242        }
243        dict.set_item("ts_event", self.ts_event.as_u64())?;
244        dict.set_item("ts_init", self.ts_init.as_u64())?;
245
246        match self.base_currency {
247            Some(base_currency) => {
248                dict.set_item("base_currency", base_currency.code.to_string())?;
249            }
250            None => dict.set_item("base_currency", "None")?,
251        }
252        Ok(dict.into())
253    }
254}