Implement event AccountState in Rust pyo3

nautilus_core/model/src/events/account/mod.rs

@@ -14,3 +14,5 @@
 // -------------------------------------------------------------------------------------------------
 
 pub mod state;
+#[cfg(feature = "stubs")]
+pub mod stubs;

nautilus_core/model/src/events/account/state.rs

@@ -13,7 +13,12 @@
 //  limitations under the License.
 // -------------------------------------------------------------------------------------------------
 
+use std::fmt::{Display, Formatter};
+
+use anyhow::Result;
 use nautilus_core::{time::UnixNanos, uuid::UUID4};
+use pyo3::prelude::*;
+use serde::{Deserialize, Serialize};
 
 use crate::{
     enums::AccountType,
@@ -25,7 +30,11 @@ use crate::{
 };
 
 #[repr(C)]
-#[derive(Debug)]
+#[derive(Debug, Serialize, Deserialize)]
+#[cfg_attr(
+    feature = "python",
+    pyclass(module = "nautilus_trader.core.nautilus_pyo3.model")
+)]
 pub struct AccountState {
     pub account_id: AccountId,
     pub account_type: AccountType,
@@ -37,3 +46,97 @@ pub struct AccountState {
     pub ts_event: UnixNanos,
     pub ts_init: UnixNanos,
 }
+
+impl AccountState {
+    #[allow(clippy::too_many_arguments)]
+    pub fn new(
+        account_id: AccountId,
+        account_type: AccountType,
+        base_currency: Currency,
+        balances: Vec<AccountBalance>,
+        margins: Vec<MarginBalance>,
+        is_reported: bool,
+        event_id: UUID4,
+        ts_event: UnixNanos,
+        ts_init: UnixNanos,
+    ) -> Result<AccountState> {
+        Ok(AccountState {
+            account_id,
+            account_type,
+            base_currency,
+            balances,
+            margins,
+            is_reported,
+            event_id,
+            ts_event,
+            ts_init,
+        })
+    }
+}
+
+impl Display for AccountState {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        write!(
+            f,
+            "AccountState(account_id={}, account_type={}, base_currency={}, is_reported={}, balances=[{}], margins=[{}], event_id={})",
+            self.account_id,
+            self.account_type,
+            self.base_currency.code,
+            self.is_reported,
+            self.balances.iter().map(|b| format!("{}", b)).collect::<Vec<String>>().join(","),
+            self.margins.iter().map(|m| format!("{}", m)).collect::<Vec<String>>().join(","),
+            self.event_id
+        )
+    }
+}
+
+impl PartialEq for AccountState {
+    fn eq(&self, other: &Self) -> bool {
+        self.account_id == other.account_id
+            && self.account_type == other.account_type
+            && self.event_id == other.event_id
+    }
+}
+
+////////////////////////////////////////////////////////////////////////////////
+// Tests
+////////////////////////////////////////////////////////////////////////////////
+#[cfg(test)]
+mod tests {
+    use rstest::rstest;
+
+    use crate::events::account::{
+        state::AccountState,
+        stubs::{cash_account_state, margin_account_state},
+    };
+
+    #[rstest]
+    fn test_equality() {
+        let cash_account_state_1 = cash_account_state();
+        let cash_account_state_2 = cash_account_state();
+        assert_eq!(cash_account_state_1, cash_account_state_2);
+    }
+
+    #[rstest]
+    fn test_display_cash_account_state(cash_account_state: AccountState) {
+        let display = format!("{}", cash_account_state);
+        assert_eq!(
+            display,
+            "AccountState(account_id=SIM-001, account_type=CASH, base_currency=USD, is_reported=true, \
+            balances=[AccountBalance(total=1525000.00 USD, locked=25000.00 USD, free=1500000.00 USD)], \
+            margins=[], event_id=16578139-a945-4b65-b46c-bc131a15d8e7)"
+        );
+    }
+
+    #[rstest]
+    fn test_display_margin_account_state(margin_account_state: AccountState) {
+        let display = format!("{}", margin_account_state);
+        assert_eq!(
+            display,
+            "AccountState(account_id=SIM-001, account_type=MARGIN, base_currency=USD, is_reported=true, \
+            balances=[AccountBalance(total=1525000.00 USD, locked=25000.00 USD, free=1500000.00 USD)], \
+            margins=[MarginBalance(initial=5000.00 USD, maintenance=20000.00 USD, instrument_id=BTCUSDT.COINBASE)], \
+            event_id=16578139-a945-4b65-b46c-bc131a15d8e7)"
+        );
+    }
+}

nautilus_core/model/src/events/account/stubs.rs

@@ -0,0 +1,58 @@
+// -------------------------------------------------------------------------------------------------
+//  Copyright (C) 2015-2024 Nautech Systems Pty Ltd. All rights reserved.
+//  https://nautechsystems.io
+//
+//  Licensed under the GNU Lesser General Public License Version 3.0 (the "License");
+//  You may not use this file except in compliance with the License.
+//  You may obtain a copy of the License at https://www.gnu.org/licenses/lgpl-3.0.en.html
+//
+//  Unless required by applicable law or agreed to in writing, software
+//  distributed under the License is distributed on an "AS IS" BASIS,
+//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+//  See the License for the specific language governing permissions and
+//  limitations under the License.
+// -------------------------------------------------------------------------------------------------
+
+use rstest::fixture;
+
+use crate::{
+    enums::AccountType,
+    events::account::state::AccountState,
+    identifiers::stubs::{account_id, uuid4},
+    types::{
+        currency::Currency,
+        stubs::{account_balance_test, margin_balance_test},
+    },
+};
+
+#[fixture]
+pub fn cash_account_state() -> AccountState {
+    AccountState::new(
+        account_id(),
+        AccountType::Cash,
+        Currency::USD(),
+        vec![account_balance_test()],
+        vec![],
+        true,
+        uuid4(),
+        0,
+        0,
+    )
+    .unwrap()
+}
+
+#[fixture]
+pub fn margin_account_state() -> AccountState {
+    AccountState::new(
+        account_id(),
+        AccountType::Margin,
+        Currency::USD(),
+        vec![account_balance_test()],
+        vec![margin_balance_test()],
+        true,
+        uuid4(),
+        0,
+        0,
+    )
+    .unwrap()
+}

nautilus_core/model/src/python/events/account/mod.rs

@@ -0,0 +1,16 @@
+// -------------------------------------------------------------------------------------------------
+//  Copyright (C) 2015-2024 Nautech Systems Pty Ltd. All rights reserved.
+//  https://nautechsystems.io
+//
+//  Licensed under the GNU Lesser General Public License Version 3.0 (the "License");
+//  You may not use this file except in compliance with the License.
+//  You may obtain a copy of the License at https://www.gnu.org/licenses/lgpl-3.0.en.html
+//
+//  Unless required by applicable law or agreed to in writing, software
+//  distributed under the License is distributed on an "AS IS" BASIS,
+//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+//  See the License for the specific language governing permissions and
+//  limitations under the License.
+// -------------------------------------------------------------------------------------------------
+
+pub mod state;

nautilus_core/model/src/python/events/account/state.rs

@@ -0,0 +1,125 @@
+// -------------------------------------------------------------------------------------------------
+//  Copyright (C) 2015-2024 Nautech Systems Pty Ltd. All rights reserved.
+//  https://nautechsystems.io
+//
+//  Licensed under the GNU Lesser General Public License Version 3.0 (the "License");
+//  You may not use this file except in compliance with the License.
+//  You may obtain a copy of the License at https://www.gnu.org/licenses/lgpl-3.0.en.html
+//
+//  Unless required by applicable law or agreed to in writing, software
+//  distributed under the License is distributed on an "AS IS" BASIS,
+//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+//  See the License for the specific language governing permissions and
+//  limitations under the License.
+// -------------------------------------------------------------------------------------------------
+
+use nautilus_core::{
+    python::{serialization::from_dict_pyo3, to_pyvalue_err},
+    time::UnixNanos,
+    uuid::UUID4,
+};
+use pyo3::{basic::CompareOp, prelude::*, types::PyDict};
+use rust_decimal::prelude::ToPrimitive;
+
+use crate::{
+    enums::AccountType,
+    events::account::state::AccountState,
+    identifiers::account_id::AccountId,
+    types::{
+        balance::{AccountBalance, MarginBalance},
+        currency::Currency,
+    },
+};
+
+#[pymethods]
+impl AccountState {
+    #[allow(clippy::too_many_arguments)]
+    #[new]
+    fn py_new(
+        account_id: AccountId,
+        account_type: AccountType,
+        base_currency: Currency,
+        balances: Vec<AccountBalance>,
+        margins: Vec<MarginBalance>,
+        is_reported: bool,
+        event_id: UUID4,
+        ts_event: UnixNanos,
+        ts_init: UnixNanos,
+    ) -> PyResult<Self> {
+        AccountState::new(
+            account_id,
+            account_type,
+            base_currency,
+            balances,
+            margins,
+            is_reported,
+            event_id,
+            ts_event,
+            ts_init,
+        )
+        .map_err(to_pyvalue_err)
+    }
+
+    fn __richcmp__(&self, other: &Self, op: CompareOp, py: Python<'_>) -> Py<PyAny> {
+        match op {
+            CompareOp::Eq => self.eq(other).into_py(py),
+            CompareOp::Ne => self.ne(other).into_py(py),
+            _ => py.NotImplemented(),
+        }
+    }
+
+    fn __repr__(&self) -> String {
+        format!(
+            "{}(account_id={},account_type={},base_currency={},balances={},margins={},is_reported={},event_id={})",
+            stringify!(AccountState),
+            self.account_id,
+            self.account_type,
+            self.base_currency.code,
+            self.balances.iter().map(|b| format!("{}", b)).collect::<Vec<String>>().join(","),
+            self.margins.iter().map(|m| format!("{}", m)).collect::<Vec<String>>().join(","),
+            self.is_reported,
+            self.event_id,
+        )
+    }
+
+    fn __str__(&self) -> String {
+        format!(
+            "{}(account_id={},account_type={},base_currency={},balances={},margins={},is_reported={},event_id={})",
+            stringify!(AccountState),
+            self.account_id,
+            self.account_type,
+            self.base_currency.code,
+            self.balances.iter().map(|b| format!("{}", b)).collect::<Vec<String>>().join(","),
+            self.margins.iter().map(|m| format!("{}", m)).collect::<Vec<String>>().join(","),
+            self.is_reported,
+            self.event_id,
+        )
+    }
+
+    #[staticmethod]
+    #[pyo3(name = "from_dict")]
+    fn py_from_dict(py: Python<'_>, values: Py<PyDict>) -> PyResult<Self> {
+        from_dict_pyo3(py, values)
+    }
+
+    #[pyo3(name = "to_dict")]
+    fn py_to_dict(&self, py: Python<'_>) -> PyResult<PyObject> {
+        let dict = PyDict::new(py);
+        dict.set_item("type", stringify!(AccountState))?;
+        dict.set_item("account_id", self.account_id.to_string())?;
+        dict.set_item("account_type", self.account_type.to_string())?;
+        dict.set_item("base_currency", self.base_currency.code.to_string())?;
+        // iterate over balances and margins and run to_dict on each item and collect them
+        let balances_dict: PyResult<Vec<_>> =
+            self.balances.iter().map(|b| b.py_to_dict(py)).collect();
+        let margins_dict: PyResult<Vec<_>> =
+            self.margins.iter().map(|m| m.py_to_dict(py)).collect();
+        dict.set_item("balances", balances_dict?)?;
+        dict.set_item("margins", margins_dict?)?;
+        dict.set_item("is_reported", self.is_reported)?;
+        dict.set_item("event_id", self.event_id.to_string())?;
+        dict.set_item("ts_event", self.ts_event.to_u64())?;
+        dict.set_item("ts_init", self.ts_init.to_u64())?;
+        Ok(dict.into())
+    }
+}

nautilus_core/model/src/python/events/mod.rs

@@ -13,4 +13,5 @@
 //  limitations under the License.
 // -------------------------------------------------------------------------------------------------
 
+pub mod account;
 pub mod order;

nautilus_core/model/src/python/mod.rs

@@ -302,7 +302,7 @@ pub fn model(_: Python<'_>, m: &PyModule) -> PyResult<()> {
     m.add_class::<crate::instruments::futures_contract::FuturesContract>()?;
     m.add_class::<crate::instruments::options_contract::OptionsContract>()?;
     m.add_class::<crate::instruments::synthetic::SyntheticInstrument>()?;
-    // Events
+    // Events - order
     m.add_class::<crate::events::order::denied::OrderDenied>()?;
     m.add_class::<crate::events::order::filled::OrderFilled>()?;
     m.add_class::<crate::events::order::initialized::OrderInitialized>()?;
@@ -319,5 +319,7 @@ pub fn model(_: Python<'_>, m: &PyModule) -> PyResult<()> {
     m.add_class::<crate::events::order::cancel_rejected::OrderCancelRejected>()?;
     m.add_class::<crate::events::order::canceled::OrderCanceled>()?;
     m.add_class::<crate::events::order::expired::OrderExpired>()?;
+    // Events - account
+    m.add_class::<crate::events::account::state::AccountState>()?;
     Ok(())
 }

nautilus_core/model/src/python/types/balance.rs

@@ -66,7 +66,7 @@ impl AccountBalance {
     }
 
     #[pyo3(name = "to_dict")]
-    fn py_to_dict(&self, py: Python<'_>) -> PyResult<PyObject> {
+    pub fn py_to_dict(&self, py: Python<'_>) -> PyResult<PyObject> {
         let dict = PyDict::new(py);
         dict.set_item("type", stringify!(AccountBalance))?;
         dict.set_item("total", self.total.to_string())?;
@@ -118,7 +118,7 @@ impl MarginBalance {
     }
 
     #[pyo3(name = "to_dict")]
-    fn py_to_dict(&self, py: Python<'_>) -> PyResult<PyObject> {
+    pub fn py_to_dict(&self, py: Python<'_>) -> PyResult<PyObject> {
         let dict = PyDict::new(py);
         dict.set_item("type", stringify!(MarginBalance))?;
         dict.set_item("initial", self.initial.to_string())?;