1use mkt_types::{ExchangeId, MarketKind};
2use strum_macros::{Display, EnumString, IntoStaticStr};
3
4#[non_exhaustive]
5#[derive(
6 Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Display, EnumString, IntoStaticStr,
7)]
8#[strum(serialize_all = "snake_case", ascii_case_insensitive)]
9pub enum Capability {
10 MarketData,
11 Account,
12 SpotTrading,
13 FuturesTrading,
14 PublicStream,
15 PrivateStream,
16}
17
18#[non_exhaustive]
19#[derive(Debug, Clone, PartialEq, Eq, Default)]
20pub struct RestCapabilities {
21 pub market_data: bool,
22 pub account: bool,
23 pub spot_trading: bool,
24 pub futures_trading: bool,
25}
26
27impl RestCapabilities {
28 pub fn with_market_data(mut self) -> Self {
29 self.market_data = true;
30 self
31 }
32
33 pub fn with_account(mut self) -> Self {
34 self.account = true;
35 self
36 }
37
38 pub fn with_spot_trading(mut self) -> Self {
39 self.spot_trading = true;
40 self
41 }
42
43 pub fn with_futures_trading(mut self) -> Self {
44 self.futures_trading = true;
45 self
46 }
47}
48
49#[non_exhaustive]
50#[derive(Debug, Clone, PartialEq, Eq, Default)]
51pub struct StreamCapabilities {
52 pub public: bool,
53 pub private: bool,
54}
55
56impl StreamCapabilities {
57 pub fn with_public(mut self) -> Self {
58 self.public = true;
59 self
60 }
61
62 pub fn with_private(mut self) -> Self {
63 self.private = true;
64 self
65 }
66}
67
68#[non_exhaustive]
69#[derive(Debug, Clone, PartialEq, Eq)]
70pub enum TransportControl {
71 OfficialSdkManaged { sdk: &'static str },
79 DirectManaged,
81}
82
83#[non_exhaustive]
84#[derive(Debug, Clone, PartialEq, Eq)]
85pub struct Capabilities {
86 pub exchange_id: ExchangeId,
87 pub markets: Vec<MarketKind>,
88 pub rest: RestCapabilities,
89 pub stream: StreamCapabilities,
90 pub transport: TransportControl,
91}
92
93impl Capabilities {
94 pub fn new(exchange_id: ExchangeId) -> Self {
95 Self {
96 exchange_id,
97 markets: Vec::new(),
98 rest: RestCapabilities::default(),
99 stream: StreamCapabilities::default(),
100 transport: TransportControl::DirectManaged,
101 }
102 }
103
104 pub fn with_markets(mut self, markets: impl IntoIterator<Item = MarketKind>) -> Self {
105 self.markets = markets.into_iter().collect();
106 self.markets.sort();
107 self.markets.dedup();
108 self
109 }
110
111 pub fn supports_market(&self, market: MarketKind) -> bool {
112 self.markets.contains(&market)
113 }
114
115 pub fn with_rest(mut self, rest: RestCapabilities) -> Self {
116 self.rest = rest;
117 self
118 }
119
120 pub fn with_stream(mut self, stream: StreamCapabilities) -> Self {
121 self.stream = stream;
122 self
123 }
124
125 pub fn with_capabilities(mut self, capabilities: impl IntoIterator<Item = Capability>) -> Self {
126 for capability in capabilities {
127 match capability {
128 Capability::MarketData => self.rest.market_data = true,
129 Capability::Account => self.rest.account = true,
130 Capability::SpotTrading => self.rest.spot_trading = true,
131 Capability::FuturesTrading => self.rest.futures_trading = true,
132 Capability::PublicStream => self.stream.public = true,
133 Capability::PrivateStream => self.stream.private = true,
134 }
135 }
136
137 self
138 }
139
140 pub fn supports_capability(&self, capability: Capability) -> bool {
141 match capability {
142 Capability::MarketData => self.rest.market_data,
143 Capability::Account => self.rest.account,
144 Capability::SpotTrading => self.rest.spot_trading,
145 Capability::FuturesTrading => self.rest.futures_trading,
146 Capability::PublicStream => self.stream.public,
147 Capability::PrivateStream => self.stream.private,
148 }
149 }
150
151 pub fn with_transport(mut self, transport: TransportControl) -> Self {
152 self.transport = transport;
153 self
154 }
155}
156
157#[cfg(test)]
158mod tests {
159 use super::{Capabilities, Capability};
160 use mkt_types::{ExchangeId, KnownExchange, MarketKind};
161
162 #[test]
163 fn market_capabilities_are_sorted_and_deduplicated() {
164 let capabilities = Capabilities::new(ExchangeId::from(KnownExchange::Binance))
165 .with_markets([
166 MarketKind::Spot,
167 MarketKind::linear_perpetual(),
168 MarketKind::Spot,
169 ]);
170
171 assert_eq!(
172 capabilities.markets,
173 vec![MarketKind::Spot, MarketKind::linear_perpetual()]
174 );
175 assert!(capabilities.supports_market(MarketKind::Spot));
176 }
177
178 #[test]
179 fn capability_queries_cover_rest_and_stream_flags() {
180 let capabilities = Capabilities::new(ExchangeId::from(KnownExchange::Binance))
181 .with_capabilities([Capability::MarketData, Capability::PrivateStream]);
182
183 assert!(capabilities.supports_capability(Capability::MarketData));
184 assert!(capabilities.supports_capability(Capability::PrivateStream));
185 assert!(!capabilities.supports_capability(Capability::Account));
186 }
187}