Skip to main content

tvdata_rs/scanner/
query.rs

1use std::collections::BTreeMap;
2
3use serde::Serialize;
4use serde_json::Value;
5
6use crate::error::{Error, Result as TvResult};
7use crate::scanner::field::{Column, Market, Ticker};
8use crate::scanner::fields::{core, fundamentals, price};
9use crate::scanner::filter::{FilterCondition, FilterTree, SortSpec};
10
11#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
12pub struct SymbolGroup {
13    #[serde(rename = "type")]
14    pub kind: String,
15    pub values: Vec<String>,
16}
17
18impl SymbolGroup {
19    pub fn new(
20        kind: impl Into<String>,
21        values: impl IntoIterator<Item = impl Into<String>>,
22    ) -> Self {
23        Self {
24            kind: kind.into(),
25            values: values.into_iter().map(Into::into).collect(),
26        }
27    }
28}
29
30#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
31pub struct Watchlist {
32    pub id: i64,
33}
34
35#[derive(Debug, Clone, PartialEq, Eq, Serialize, Default)]
36pub struct SymbolQuery {
37    pub types: Vec<String>,
38}
39
40#[derive(Debug, Clone, PartialEq, Eq, Serialize, Default)]
41pub struct Symbols {
42    pub query: SymbolQuery,
43    pub tickers: Vec<Ticker>,
44    #[serde(skip_serializing_if = "Vec::is_empty", default)]
45    pub symbolset: Vec<String>,
46    #[serde(skip_serializing_if = "Option::is_none")]
47    pub watchlist: Option<Watchlist>,
48    #[serde(skip_serializing_if = "Vec::is_empty", default)]
49    pub groups: Vec<SymbolGroup>,
50}
51
52impl Symbols {
53    pub fn with_tickers<I, T>(mut self, tickers: I) -> Self
54    where
55        I: IntoIterator<Item = T>,
56        T: Into<Ticker>,
57    {
58        self.tickers = tickers.into_iter().map(Into::into).collect();
59        self
60    }
61
62    pub fn with_symbolset<I, S>(mut self, symbolset: I) -> Self
63    where
64        I: IntoIterator<Item = S>,
65        S: Into<String>,
66    {
67        self.symbolset = symbolset.into_iter().map(Into::into).collect();
68        self
69    }
70
71    pub fn with_watchlist(mut self, id: i64) -> Self {
72        self.watchlist = Some(Watchlist { id });
73        self
74    }
75
76    pub fn with_group(mut self, group: SymbolGroup) -> Self {
77        self.groups.push(group);
78        self
79    }
80
81    pub fn with_types<I, S>(mut self, types: I) -> Self
82    where
83        I: IntoIterator<Item = S>,
84        S: Into<String>,
85    {
86        self.query.types = types.into_iter().map(Into::into).collect();
87        self
88    }
89}
90
91#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize)]
92pub struct Page([usize; 2]);
93
94impl Page {
95    pub fn new(offset: usize, limit: usize) -> TvResult<Self> {
96        if limit == 0 {
97            return Err(Error::InvalidPageLimit);
98        }
99        Ok(Self([offset, limit]))
100    }
101}
102
103impl Default for Page {
104    fn default() -> Self {
105        Self([0, 50])
106    }
107}
108
109#[derive(Debug, Clone, PartialEq, Eq)]
110pub enum PriceConversion {
111    SymbolCurrency,
112    MarketCurrency,
113    Specific(String),
114}
115
116impl Serialize for PriceConversion {
117    fn serialize<S>(&self, serializer: S) -> std::result::Result<S::Ok, S::Error>
118    where
119        S: serde::Serializer,
120    {
121        use serde::ser::SerializeMap;
122
123        let mut map = serializer.serialize_map(Some(1))?;
124        match self {
125            Self::SymbolCurrency => map.serialize_entry("to_symbol", &true)?,
126            Self::MarketCurrency => map.serialize_entry("to_symbol", &false)?,
127            Self::Specific(currency) => {
128                map.serialize_entry("to_currency", &currency.to_lowercase())?
129            }
130        }
131        map.end()
132    }
133}
134
135#[derive(Debug, Clone, PartialEq, Serialize)]
136pub struct ScanQuery {
137    #[serde(skip_serializing_if = "Vec::is_empty", default)]
138    pub markets: Vec<Market>,
139    pub symbols: Symbols,
140    #[serde(skip_serializing_if = "BTreeMap::is_empty", default)]
141    pub options: BTreeMap<String, Value>,
142    pub columns: Vec<Column>,
143    #[serde(rename = "filter", skip_serializing_if = "Vec::is_empty", default)]
144    pub filters: Vec<FilterCondition>,
145    #[serde(rename = "filter2", skip_serializing_if = "Option::is_none")]
146    pub filter_tree: Option<FilterTree>,
147    #[serde(skip_serializing_if = "Option::is_none")]
148    pub sort: Option<SortSpec>,
149    #[serde(rename = "range")]
150    pub page: Page,
151    #[serde(skip_serializing_if = "Option::is_none")]
152    pub ignore_unknown_fields: Option<bool>,
153    #[serde(skip_serializing_if = "Option::is_none")]
154    pub preset: Option<String>,
155    #[serde(skip_serializing_if = "Option::is_none")]
156    pub price_conversion: Option<PriceConversion>,
157}
158
159impl Default for ScanQuery {
160    fn default() -> Self {
161        Self {
162            markets: Vec::new(),
163            symbols: Symbols::default(),
164            options: BTreeMap::from([(String::from("lang"), Value::String(String::from("en")))]),
165            columns: vec![
166                core::NAME,
167                price::CLOSE,
168                price::VOLUME,
169                fundamentals::MARKET_CAP_BASIC,
170            ],
171            filters: Vec::new(),
172            filter_tree: None,
173            sort: None,
174            page: Page::default(),
175            ignore_unknown_fields: None,
176            preset: None,
177            price_conversion: None,
178        }
179    }
180}
181
182impl ScanQuery {
183    pub fn new() -> Self {
184        Self::default()
185    }
186
187    pub fn select<I, C>(mut self, columns: I) -> Self
188    where
189        I: IntoIterator<Item = C>,
190        C: Into<Column>,
191    {
192        self.columns = columns.into_iter().map(Into::into).collect();
193        self
194    }
195
196    pub fn push_column(mut self, column: impl Into<Column>) -> Self {
197        self.columns.push(column.into());
198        self
199    }
200
201    pub fn market(mut self, market: impl Into<Market>) -> Self {
202        self.markets = vec![market.into()];
203        self
204    }
205
206    pub fn markets<I, M>(mut self, markets: I) -> Self
207    where
208        I: IntoIterator<Item = M>,
209        M: Into<Market>,
210    {
211        self.markets = markets.into_iter().map(Into::into).collect();
212        self
213    }
214
215    pub fn symbols(mut self, symbols: Symbols) -> Self {
216        self.symbols = symbols;
217        self
218    }
219
220    pub fn tickers<I, T>(mut self, tickers: I) -> Self
221    where
222        I: IntoIterator<Item = T>,
223        T: Into<Ticker>,
224    {
225        self.symbols = self.symbols.with_tickers(tickers);
226        self
227    }
228
229    pub fn symbolset<I, S>(mut self, symbolset: I) -> Self
230    where
231        I: IntoIterator<Item = S>,
232        S: Into<String>,
233    {
234        self.symbols = self.symbols.with_symbolset(symbolset);
235        self
236    }
237
238    pub fn watchlist(mut self, id: i64) -> Self {
239        self.symbols = self.symbols.with_watchlist(id);
240        self
241    }
242
243    pub fn group(mut self, group: SymbolGroup) -> Self {
244        self.symbols = self.symbols.with_group(group);
245        self
246    }
247
248    pub fn symbol_types<I, S>(mut self, types: I) -> Self
249    where
250        I: IntoIterator<Item = S>,
251        S: Into<String>,
252    {
253        self.symbols = self.symbols.with_types(types);
254        self
255    }
256
257    pub fn filter(mut self, filter: FilterCondition) -> Self {
258        self.filters.push(filter);
259        self
260    }
261
262    pub fn filters<I>(mut self, filters: I) -> Self
263    where
264        I: IntoIterator<Item = FilterCondition>,
265    {
266        self.filters.extend(filters);
267        self
268    }
269
270    pub fn filter_tree(mut self, filter_tree: FilterTree) -> Self {
271        self.filter_tree = Some(filter_tree);
272        self
273    }
274
275    pub fn sort(mut self, sort: SortSpec) -> Self {
276        self.sort = Some(sort);
277        self
278    }
279
280    pub fn page(mut self, offset: usize, limit: usize) -> TvResult<Self> {
281        self.page = Page::new(offset, limit)?;
282        Ok(self)
283    }
284
285    pub fn language(mut self, language: impl Into<String>) -> Self {
286        self.options
287            .insert(String::from("lang"), Value::String(language.into()));
288        self
289    }
290
291    pub fn option(mut self, key: impl Into<String>, value: impl Into<Value>) -> Self {
292        self.options.insert(key.into(), value.into());
293        self
294    }
295
296    pub fn preset(mut self, preset: impl Into<String>) -> Self {
297        self.preset = Some(preset.into());
298        self
299    }
300
301    pub fn price_conversion(mut self, price_conversion: PriceConversion) -> Self {
302        self.price_conversion = Some(price_conversion);
303        self
304    }
305
306    pub fn ignore_unknown_fields(mut self, ignore_unknown_fields: bool) -> Self {
307        self.ignore_unknown_fields = Some(ignore_unknown_fields);
308        self
309    }
310
311    pub fn route_segment(&self) -> String {
312        let requires_global_route = !self.symbols.symbolset.is_empty()
313            || self.symbols.watchlist.is_some()
314            || !self.symbols.groups.is_empty();
315
316        match (requires_global_route, self.markets.as_slice()) {
317            (false, [market]) => format!("{}/scan", market.as_str()),
318            _ => String::from("global/scan"),
319        }
320    }
321}
322
323#[cfg(test)]
324mod tests {
325    use serde_json::json;
326
327    use super::*;
328    use crate::scanner::fields::{analyst, technical};
329
330    #[test]
331    fn uses_global_route_for_multi_market_or_symbolset_queries() {
332        let query = ScanQuery::new()
333            .markets(["america", "crypto"])
334            .select([core::NAME, price::CLOSE]);
335        assert_eq!(query.route_segment(), "global/scan");
336
337        let query = ScanQuery::new()
338            .symbolset(["SYML:SP;SPX"])
339            .preset("index_components_market_pages");
340        assert_eq!(query.route_segment(), "global/scan");
341
342        let query = ScanQuery::new()
343            .market("america")
344            .symbolset(["SYML:SP;SPX"])
345            .preset("index_components_market_pages");
346        assert_eq!(query.route_segment(), "global/scan");
347    }
348
349    #[test]
350    fn uses_global_route_for_watchlist_and_group_queries() {
351        let query = ScanQuery::new().market("america").watchlist(42);
352        assert_eq!(query.route_segment(), "global/scan");
353
354        let query = ScanQuery::new()
355            .market("america")
356            .group(SymbolGroup::new("index", ["SPX"]));
357        assert_eq!(query.route_segment(), "global/scan");
358    }
359
360    #[test]
361    fn serializes_tradingview_scan_payload() {
362        let query = ScanQuery::new()
363            .market("america")
364            .tickers(["NASDAQ:AAPL"])
365            .select([
366                core::NAME,
367                price::CLOSE,
368                analyst::PRICE_TARGET_AVERAGE,
369                technical::RSI.with_interval("1W"),
370            ])
371            .filter(price::CLOSE.clone().gt(100))
372            .sort(price::CLOSE.clone().sort(crate::scanner::SortOrder::Desc));
373
374        let value = serde_json::to_value(query).unwrap();
375        assert_eq!(value["columns"][0], "name");
376        assert_eq!(value["columns"][3], "RSI|1W");
377        assert_eq!(value["filter"][0]["operation"], "greater");
378        assert_eq!(value["markets"], json!(["america"]));
379    }
380}