diff --git a/vega_sim/api/data.py b/vega_sim/api/data.py index a4e5ff1ca..989e48525 100644 --- a/vega_sim/api/data.py +++ b/vega_sim/api/data.py @@ -40,7 +40,7 @@ class MissingMarketError(Exception): PriceLevel = namedtuple("PriceLevel", ["price", "number_of_orders", "volume"]) -@dataclass +@dataclass(frozen=True) class AccountData: owner: str balance: float @@ -53,14 +53,14 @@ def account_id(self): return f"{self.owner}-{self.type}-{self.market_id}-{self.asset}" -@dataclass +@dataclass(frozen=True) class IcebergOrder: peak_size: float minimum_visible_size: float reserved_remaining: float -@dataclass +@dataclass(frozen=True) class Order: price: float size: float @@ -127,20 +127,20 @@ class Order: ) -@dataclass +@dataclass(frozen=True) class LiquidationPrice: open_volume_only: float including_buy_orders: float including_sell_orders: float -@dataclass +@dataclass(frozen=True) class MarginEstimate: best_case: MarginLevels worst_case: MarginLevels -@dataclass +@dataclass(frozen=True) class LiquidationEstimate: best_case: LiquidationPrice worst_case: LiquidationPrice @@ -153,7 +153,7 @@ class DecimalSpec: asset_decimals: Optional[int] = None -@dataclass +@dataclass(frozen=True) class LedgerEntry: from_account: vega_protos.vega.AccountDetails to_account: vega_protos.vega.AccountDetails @@ -162,7 +162,7 @@ class LedgerEntry: timestamp: int -@dataclass +@dataclass(frozen=True) class AggregatedLedgerEntry: timestamp: int quantity: float @@ -176,19 +176,19 @@ class AggregatedLedgerEntry: to_account_market_id: str -@dataclass +@dataclass(frozen=True) class MarketDepth: buys: List[PriceLevel] sells: List[PriceLevel] -@dataclass +@dataclass(frozen=True) class OrdersBySide: bids: List[Order] asks: List[Order] -@dataclass +@dataclass(frozen=True) class Fee: maker_fee: float infrastructure_fee: float diff --git a/vega_sim/local_data_cache.py b/vega_sim/local_data_cache.py index b33beec4d..1af29813f 100644 --- a/vega_sim/local_data_cache.py +++ b/vega_sim/local_data_cache.py @@ -599,32 +599,37 @@ def get_trades_from_stream( def get_accounts_from_stream( self, market_id: Optional[str] = None, - asset_id: Optional[str] = None, party_id: Optional[str] = None, + asset_id: Optional[str] = None, ) -> List[data.AccountData]: """Loads accounts for either a given party, market or both from stream. - Must pass one or the other Args: market_id: - optional str, Restrict to trades on a specific market + optional str, Restrict to accounts on a specific market party_id: - optional str, Select only trades with a given id + optional str, Select only accounts with a given party id + asset_id: + optional str, Restrict to only accounts for a given asset Returns: List[AccountData], list of formatted trade objects which match the required restrictions. """ - if market_id is None and party_id is None: - raise Exception("At least one of market_id and party_id must be specified") + with self.account_lock: - to_check_keys = set() - if market_id is not None: - to_check_keys.update( - self._account_keys_for_market.get(market_id, set()) - ) - if party_id is not None: - to_check_keys.update(self._account_keys_for_party.get(party_id, set())) + if market_id is None and party_id is None: + to_check_keys = self._accounts_from_feed.keys() + else: + to_check_keys = set() + if market_id is not None: + to_check_keys.update( + self._account_keys_for_market.get(market_id, set()) + ) + if party_id is not None: + to_check_keys.update( + self._account_keys_for_party.get(party_id, set()) + ) results = [] for key in to_check_keys: acct = self._accounts_from_feed[key] @@ -643,3 +648,17 @@ def get_accounts_from_stream( results.append(acct) return results + + def get_all_accounts_from_stream( + self, + ) -> List[data.AccountData]: + """Loads all accounts stored. + + Returns: + List[AccountData], list of formatted account objects + """ + + with self.account_lock: + accts = list(self._accounts_from_feed.values()) + + return accts diff --git a/vega_sim/scenario/common/agents.py b/vega_sim/scenario/common/agents.py index cb168d5bd..93dc7d96f 100644 --- a/vega_sim/scenario/common/agents.py +++ b/vega_sim/scenario/common/agents.py @@ -2865,7 +2865,7 @@ def step(self, vega_state: VegaState): self.seen_trades.add(trade.id) market_trades.setdefault(market.id, []).append(trade) - accounts = self.vega.list_accounts() + accounts = self.vega.get_accounts_from_stream() self.states.append( MarketHistoryData( at_time=start_time,