tvscreener/core/
stock.rs

1// Copyright 2026 tvscreener-rs contributors
2// SPDX-License-Identifier: Apache-2.0
3
4//! [`StockScreener`] (`TradingView` subtype `global`).
5//!
6//! Default stock fields, markets, and symbol-type filters.
7
8use super::Screener;
9use crate::error::{Result, TvscreenerError};
10use crate::field::{
11    default_stock_fields, stock_market_capitalization, symbol_type, type_filter, Market,
12    SymbolType, TypeFilter,
13};
14use crate::filter::FilterOperator;
15use serde_json::json;
16use std::collections::HashSet;
17
18const SYMBOL_TYPE_TO_TYPE: &[(&str, &str)] = &[
19    ("COMMON_STOCK", "STOCK"),
20    ("DEPOSITORY_RECEIPT", "DEPOSITORY_RECEIPT"),
21    ("ETF", "FUND"),
22    ("MUTUAL_FUND", "FUND"),
23    ("REIT", "FUND"),
24    ("PREFERRED_STOCK", "STOCK"),
25    ("ETN", "STRUCTURED"),
26    ("STRUCTURED", "STRUCTURED"),
27    ("UIT", "FUND"),
28];
29
30/// Stock screener with asset-class defaults
31#[derive(Debug, Clone)]
32pub struct StockScreener {
33    inner: Screener,
34}
35
36impl StockScreener {
37    /// Creates a screener with default stock fields, America market, and market-cap sort.
38    pub fn new() -> Self {
39        let mut inner = Screener::new("global");
40        inner.set_url(crate::util::get_url("global"));
41        inner.set_all_fields_fn(default_stock_fields);
42        inner.select(default_stock_fields());
43        inner.set_markets([Market::america().value]);
44        inner.sort_by(stock_market_capitalization().field_name, false);
45        Self { inner }
46    }
47
48    /// Sets market filters
49    pub fn set_markets(&mut self, markets: impl IntoIterator<Item = Market>) -> &mut Self {
50        let selected: Vec<Market> = markets.into_iter().collect();
51        let wire_values: Vec<String> = if selected.iter().any(|m| m.const_name == "ALL") {
52            crate::field::all_markets()
53                .iter()
54                .map(|m| m.value.clone())
55                .collect()
56        } else {
57            selected.into_iter().map(|m| m.value).collect()
58        };
59        self.inner.set_markets(wire_values);
60        self
61    }
62
63    /// Adds stock type and subtype filters
64    ///
65    /// # Errors
66    ///
67    /// Returns [`TvscreenerError::InvalidRequest`] when a symbol type is unknown or
68    /// has no type mapping.
69    pub fn set_symbol_types(
70        &mut self,
71        symbol_types: impl IntoIterator<Item = impl AsRef<str>>,
72    ) -> Result<&mut Self> {
73        let mut requested = Vec::new();
74        for name in symbol_types {
75            let const_name = name.as_ref();
76            let st = symbol_type(const_name).ok_or_else(|| {
77                TvscreenerError::InvalidRequest(format!("unknown symbol type: {const_name}"))
78            })?;
79            requested.push(st);
80        }
81
82        let type_values = collect_type_filters(&requested)?;
83        emit_type_filters(&mut self.inner, &type_values);
84        emit_subtype_filters(&mut self.inner, requested);
85
86        Ok(self)
87    }
88}
89
90fn collect_type_filters(requested: &[SymbolType]) -> Result<Vec<TypeFilter>> {
91    let mut types_to_add = HashSet::new();
92    for st in requested {
93        let mapped = SYMBOL_TYPE_TO_TYPE
94            .iter()
95            .find(|(sym, _)| *sym == st.const_name)
96            .map(|(_, ty)| *ty)
97            .ok_or_else(|| {
98                TvscreenerError::InvalidRequest(format!(
99                    "symbol type `{}` has no type mapping",
100                    st.const_name
101                ))
102            })?;
103        types_to_add.insert(mapped.to_string());
104    }
105    Ok(types_to_add
106        .into_iter()
107        .filter_map(|name| type_filter(&name))
108        .collect())
109}
110
111fn emit_type_filters(inner: &mut Screener, type_values: &[TypeFilter]) {
112    if type_values.len() > 1 {
113        inner
114            .add_filter(
115                "type",
116                FilterOperator::InRange,
117                type_values.iter().map(|t| json!(t.value)),
118            )
119            .expect("type is always an allowed filter left");
120    } else if let Some(one) = type_values.first() {
121        inner
122            .add_filter("type", FilterOperator::Equal, [json!(one.value)])
123            .expect("type is always an allowed filter left");
124    }
125}
126
127fn emit_subtype_filters(inner: &mut Screener, mut symbol_types_to_apply: Vec<SymbolType>) {
128    if symbol_types_to_apply
129        .iter()
130        .any(|st| st.const_name == "COMMON_STOCK")
131        && !symbol_types_to_apply
132            .iter()
133            .any(|st| st.const_name == "DEPOSITORY_RECEIPT")
134    {
135        if let Some(dr) = symbol_type("DEPOSITORY_RECEIPT") {
136            symbol_types_to_apply.push(dr);
137        }
138    }
139
140    for st in symbol_types_to_apply {
141        inner
142            .add_filter(
143                "subtype",
144                FilterOperator::InRange,
145                st.values.iter().map(|value| json!(value)),
146            )
147            .expect("subtype is always an allowed filter left");
148    }
149}
150
151impl_typed_screener_common!(StockScreener);
152
153#[cfg(test)]
154mod tests {
155    use super::*;
156    use serde_json::{json, Value};
157
158    #[test]
159    fn set_symbol_types_common_stock_adds_depository_receipt() {
160        let mut ss = StockScreener::new();
161        ss.set_symbol_types(["COMMON_STOCK"]).expect("types");
162        let payload = ss.inner().build_payload().expect("payload");
163        let filters = payload["filter"].as_array().expect("filter");
164        let subtype_rights: Vec<&Value> = filters
165            .iter()
166            .filter(|f| f["left"] == "subtype")
167            .flat_map(|f| f["right"].as_array().into_iter().flatten())
168            .collect();
169        assert!(
170            subtype_rights.len() >= 2,
171            "COMMON_STOCK should also emit DEPOSITORY_RECEIPT subtypes; got {subtype_rights:?}"
172        );
173    }
174
175    #[test]
176    fn set_symbol_types_unknown_is_invalid() {
177        let mut ss = StockScreener::new();
178        let err = ss
179            .set_symbol_types(["NOT_A_TYPE"])
180            .expect_err("unknown type");
181        assert!(matches!(err, TvscreenerError::InvalidRequest(_)));
182    }
183
184    #[test]
185    fn set_symbol_types_multi_promotes_type_to_in_range() {
186        let mut ss = StockScreener::new();
187        ss.set_symbol_types(["ETF", "COMMON_STOCK"]).expect("types");
188        let payload = ss.inner().build_payload().expect("payload");
189        let filters = payload["filter"].as_array().expect("filter");
190        let type_filter = filters
191            .iter()
192            .find(|f| f["left"] == "type")
193            .expect("type filter");
194        assert_eq!(type_filter["operation"], json!("in_range"));
195        assert!(type_filter["right"].as_array().unwrap().len() >= 2);
196    }
197}
tvscreener/core/stock.rs

tvscreener/core/
stock.rs
