Skip to main content

tvdata_rs/search/
mod.rs

1use std::collections::BTreeMap;
2
3use bon::Builder;
4use serde::Deserialize;
5use serde_json::Value;
6
7#[cfg(test)]
8mod tests;
9
10fn default_search_language() -> String {
11    "en".to_owned()
12}
13
14fn default_search_domain() -> String {
15    "production".to_owned()
16}
17
18#[derive(Debug, Clone, Copy, PartialEq, Eq)]
19pub enum SearchAssetClass {
20    Equity,
21    Forex,
22    Crypto,
23    Futures,
24    Index,
25    Bond,
26    Cfd,
27    Option,
28}
29
30impl SearchAssetClass {
31    pub const fn api_search_type(self) -> Option<&'static str> {
32        match self {
33            Self::Equity => Some("stock"),
34            Self::Forex => Some("forex"),
35            Self::Crypto => Some("crypto"),
36            Self::Futures => Some("futures"),
37            Self::Index => Some("index"),
38            Self::Bond => Some("bond"),
39            Self::Cfd => Some("cfd"),
40            Self::Option => None,
41        }
42    }
43}
44
45#[derive(Debug, Clone, PartialEq, Eq, Builder)]
46pub struct SearchRequest {
47    #[builder(into)]
48    pub text: String,
49    #[builder(into)]
50    pub exchange: Option<String>,
51    #[builder(into)]
52    pub instrument_type: Option<String>,
53    #[builder(default)]
54    pub start: usize,
55    #[builder(default = true)]
56    pub highlight: bool,
57    #[builder(default = default_search_language(), into)]
58    pub language: String,
59    #[builder(default = default_search_domain(), into)]
60    pub domain: String,
61}
62
63impl SearchRequest {
64    pub fn new(text: impl Into<String>) -> Self {
65        Self::builder().text(text).build()
66    }
67
68    pub fn equities(text: impl Into<String>) -> Self {
69        Self::new(text).asset_class(SearchAssetClass::Equity)
70    }
71
72    pub fn forex(text: impl Into<String>) -> Self {
73        Self::new(text).asset_class(SearchAssetClass::Forex)
74    }
75
76    pub fn crypto(text: impl Into<String>) -> Self {
77        Self::new(text).asset_class(SearchAssetClass::Crypto)
78    }
79
80    pub fn futures(text: impl Into<String>) -> Self {
81        Self::new(text).asset_class(SearchAssetClass::Futures)
82    }
83
84    pub fn indices(text: impl Into<String>) -> Self {
85        Self::new(text).asset_class(SearchAssetClass::Index)
86    }
87
88    pub fn bonds(text: impl Into<String>) -> Self {
89        Self::new(text).asset_class(SearchAssetClass::Bond)
90    }
91
92    pub fn cfds(text: impl Into<String>) -> Self {
93        Self::new(text).asset_class(SearchAssetClass::Cfd)
94    }
95
96    pub fn options(text: impl Into<String>) -> Self {
97        Self::new(text)
98    }
99
100    pub fn exchange(mut self, exchange: impl Into<String>) -> Self {
101        self.exchange = Some(exchange.into());
102        self
103    }
104
105    pub fn asset_class(mut self, asset_class: SearchAssetClass) -> Self {
106        self.instrument_type = asset_class.api_search_type().map(str::to_owned);
107        self
108    }
109
110    pub fn instrument_type(mut self, instrument_type: impl Into<String>) -> Self {
111        self.instrument_type = Some(instrument_type.into());
112        self
113    }
114
115    pub fn start(mut self, start: usize) -> Self {
116        self.start = start;
117        self
118    }
119
120    pub fn highlight(mut self, highlight: bool) -> Self {
121        self.highlight = highlight;
122        self
123    }
124
125    pub fn language(mut self, language: impl Into<String>) -> Self {
126        self.language = language.into();
127        self
128    }
129
130    pub fn domain(mut self, domain: impl Into<String>) -> Self {
131        self.domain = domain.into();
132        self
133    }
134
135    pub(crate) fn to_query_pairs(&self) -> Vec<(&str, String)> {
136        let mut pairs = vec![
137            ("text", self.text.clone()),
138            ("start", self.start.to_string()),
139            ("hl", if self.highlight { "1" } else { "0" }.to_owned()),
140            ("lang", self.language.clone()),
141            ("domain", self.domain.clone()),
142        ];
143
144        if let Some(exchange) = self.exchange.as_ref().filter(|value| !value.is_empty()) {
145            pairs.push(("exchange", exchange.clone()));
146        }
147
148        if let Some(instrument_type) = self
149            .instrument_type
150            .as_ref()
151            .filter(|value| !value.is_empty())
152        {
153            pairs.push(("search_type", instrument_type.clone()));
154        }
155
156        pairs
157    }
158}
159
160#[derive(Debug, Clone, PartialEq)]
161pub struct SearchResponse {
162    pub hits: Vec<SearchHit>,
163    pub symbols_remaining: usize,
164}
165
166impl SearchResponse {
167    pub fn filtered<F>(self, predicate: F) -> Self
168    where
169        F: FnMut(&SearchHit) -> bool,
170    {
171        Self {
172            symbols_remaining: self.symbols_remaining,
173            hits: self.hits.into_iter().filter(predicate).collect(),
174        }
175    }
176}
177
178#[derive(Debug, Clone, PartialEq)]
179pub struct SearchHit {
180    pub symbol: String,
181    pub highlighted_symbol: Option<String>,
182    pub description: Option<String>,
183    pub highlighted_description: Option<String>,
184    pub instrument_type: Option<String>,
185    pub exchange: Option<String>,
186    pub country: Option<String>,
187    pub currency_code: Option<String>,
188    pub currency_logoid: Option<String>,
189    pub provider_id: Option<String>,
190    pub source_id: Option<String>,
191    pub cik_code: Option<String>,
192    pub isin: Option<String>,
193    pub cusip: Option<String>,
194    pub found_by_isin: Option<bool>,
195    pub found_by_cusip: Option<bool>,
196    pub is_primary_listing: Option<bool>,
197    pub logoid: Option<String>,
198    pub logo: Option<SearchLogo>,
199    pub source: Option<SearchSource>,
200    pub type_specs: Vec<String>,
201    pub extra: BTreeMap<String, Value>,
202}
203
204impl SearchHit {
205    pub fn is_option_like(&self) -> bool {
206        matches!(self.instrument_type.as_deref(), Some("option"))
207            || self
208                .type_specs
209                .iter()
210                .any(|value| value.eq_ignore_ascii_case("option"))
211            || self.extra.contains_key("option-type")
212            || self.extra.contains_key("expiration")
213            || self.extra.contains_key("strike")
214    }
215}
216
217#[derive(Debug, Clone, PartialEq, Eq, Deserialize)]
218pub struct SearchLogo {
219    pub style: Option<String>,
220    pub logoid: Option<String>,
221}
222
223#[derive(Debug, Clone, PartialEq, Eq, Deserialize)]
224pub struct SearchSource {
225    pub id: Option<String>,
226    pub name: Option<String>,
227    pub description: Option<String>,
228}
229
230#[derive(Debug, Clone, Deserialize)]
231pub(crate) struct RawSearchHit {
232    pub symbol: String,
233    #[serde(default)]
234    pub description: Option<String>,
235    #[serde(rename = "type", default)]
236    pub instrument_type: Option<String>,
237    #[serde(default)]
238    pub exchange: Option<String>,
239    #[serde(default)]
240    pub country: Option<String>,
241    #[serde(default)]
242    pub currency_code: Option<String>,
243    #[serde(default, rename = "currency-logoid")]
244    pub currency_logoid: Option<String>,
245    #[serde(default)]
246    pub provider_id: Option<String>,
247    #[serde(default)]
248    pub source_id: Option<String>,
249    #[serde(default)]
250    pub cik_code: Option<String>,
251    #[serde(default)]
252    pub isin: Option<String>,
253    #[serde(default)]
254    pub cusip: Option<String>,
255    #[serde(default)]
256    pub found_by_isin: Option<bool>,
257    #[serde(default)]
258    pub found_by_cusip: Option<bool>,
259    #[serde(default)]
260    pub is_primary_listing: Option<bool>,
261    #[serde(default)]
262    pub logoid: Option<String>,
263    #[serde(default)]
264    pub logo: Option<SearchLogo>,
265    #[serde(default, rename = "source2")]
266    pub source: Option<SearchSource>,
267    #[serde(default, rename = "typespecs")]
268    pub type_specs: Vec<String>,
269    #[serde(flatten)]
270    pub extra: BTreeMap<String, Value>,
271}
272
273#[derive(Debug, Clone, Deserialize)]
274pub(crate) struct RawSearchResponseV3 {
275    #[serde(default)]
276    pub symbols_remaining: usize,
277    #[serde(default, rename = "symbols")]
278    pub hits: Vec<RawSearchHit>,
279}
280
281#[derive(Debug, Clone, Deserialize)]
282#[serde(untagged)]
283pub(crate) enum RawSearchResponse {
284    V3(RawSearchResponseV3),
285    Legacy(Vec<RawSearchHit>),
286}
287
288pub(crate) fn sanitize_response(raw_response: RawSearchResponse) -> SearchResponse {
289    match raw_response {
290        RawSearchResponse::V3(response) => SearchResponse {
291            hits: sanitize_hits(response.hits),
292            symbols_remaining: response.symbols_remaining,
293        },
294        RawSearchResponse::Legacy(hits) => SearchResponse {
295            symbols_remaining: 0,
296            hits: sanitize_hits(hits),
297        },
298    }
299}
300
301pub(crate) fn sanitize_hits(raw_hits: Vec<RawSearchHit>) -> Vec<SearchHit> {
302    raw_hits.into_iter().map(sanitize_hit).collect()
303}
304
305fn sanitize_hit(hit: RawSearchHit) -> SearchHit {
306    SearchHit {
307        symbol: strip_em_tags(&hit.symbol),
308        highlighted_symbol: contains_highlight_markup(&hit.symbol).then_some(hit.symbol),
309        description: hit.description.as_ref().map(|value| strip_em_tags(value)),
310        highlighted_description: hit
311            .description
312            .filter(|value| contains_highlight_markup(value)),
313        instrument_type: hit.instrument_type,
314        exchange: hit.exchange,
315        country: hit.country,
316        currency_code: hit.currency_code,
317        currency_logoid: hit.currency_logoid,
318        provider_id: hit.provider_id,
319        source_id: hit.source_id,
320        cik_code: hit.cik_code,
321        isin: hit.isin,
322        cusip: hit.cusip,
323        found_by_isin: hit.found_by_isin,
324        found_by_cusip: hit.found_by_cusip,
325        is_primary_listing: hit.is_primary_listing,
326        logoid: hit.logoid,
327        logo: hit.logo,
328        source: hit.source,
329        type_specs: hit.type_specs,
330        extra: hit.extra,
331    }
332}
333
334fn contains_highlight_markup(value: &str) -> bool {
335    value.contains("<em>") || value.contains("</em>")
336}
337
338fn strip_em_tags(value: &str) -> String {
339    value.replace("<em>", "").replace("</em>", "")
340}