Skip to main content

rust_tg_bot_ext/utils/
update_parsing.rs

1//! Helper functions for parsing update-related inputs.
2//!
3//! Port of `telegram.ext._utils._update_parsing`.
4//! These are library-internal utilities for normalising handler filter inputs.
5
6use std::collections::HashSet;
7
8/// Accept a single chat ID or a collection of chat IDs and return a
9/// `HashSet<i64>`.
10///
11/// If `None` is passed, an empty set is returned.
12pub fn parse_chat_id(chat_id: Option<SingleOrCollection<i64>>) -> HashSet<i64> {
13    match chat_id {
14        None => HashSet::new(),
15        Some(SingleOrCollection::Single(id)) => {
16            let mut set = HashSet::with_capacity(1);
17            set.insert(id);
18            set
19        }
20        Some(SingleOrCollection::Collection(ids)) => ids.into_iter().collect(),
21    }
22}
23
24/// Accept a single username or a collection of usernames and return a
25/// `HashSet<String>` with the leading `@` stripped.
26///
27/// If `None` is passed, an empty set is returned.
28pub fn parse_username(username: Option<SingleOrCollection<String>>) -> HashSet<String> {
29    match username {
30        None => HashSet::new(),
31        Some(SingleOrCollection::Single(u)) => {
32            let mut set = HashSet::with_capacity(1);
33            set.insert(strip_at(u));
34            set
35        }
36        Some(SingleOrCollection::Collection(us)) => us.into_iter().map(strip_at).collect(),
37    }
38}
39
40/// Mirrors Python's `SCT[T]` (Single-or-Collection-of-T).
41#[derive(Debug, Clone)]
42#[non_exhaustive]
43pub enum SingleOrCollection<T> {
44    /// A single value.
45    Single(T),
46    /// A collection of values.
47    Collection(Vec<T>),
48}
49
50impl<T> From<T> for SingleOrCollection<T> {
51    fn from(value: T) -> Self {
52        Self::Single(value)
53    }
54}
55
56impl<T> From<Vec<T>> for SingleOrCollection<T> {
57    fn from(value: Vec<T>) -> Self {
58        Self::Collection(value)
59    }
60}
61
62/// Strip a leading `@` from a username string.
63fn strip_at(mut s: String) -> String {
64    if s.starts_with('@') {
65        s.remove(0);
66    }
67    s
68}
69
70#[cfg(test)]
71mod tests {
72    use super::*;
73
74    #[test]
75    fn parse_chat_id_none() {
76        assert!(parse_chat_id(None).is_empty());
77    }
78
79    #[test]
80    fn parse_chat_id_single() {
81        let result = parse_chat_id(Some(42i64.into()));
82        assert_eq!(result.len(), 1);
83        assert!(result.contains(&42));
84    }
85
86    #[test]
87    fn parse_chat_id_collection() {
88        let result = parse_chat_id(Some(vec![1i64, 2, 3].into()));
89        assert_eq!(result.len(), 3);
90        assert!(result.contains(&1));
91        assert!(result.contains(&2));
92        assert!(result.contains(&3));
93    }
94
95    #[test]
96    fn parse_username_strips_at() {
97        let result = parse_username(Some("@testuser".to_owned().into()));
98        assert!(result.contains("testuser"));
99    }
100
101    #[test]
102    fn parse_username_no_at() {
103        let result = parse_username(Some("testuser".to_owned().into()));
104        assert!(result.contains("testuser"));
105    }
106
107    #[test]
108    fn parse_username_collection() {
109        let input = vec!["@alice".to_owned(), "bob".to_owned()];
110        let result = parse_username(Some(input.into()));
111        assert!(result.contains("alice"));
112        assert!(result.contains("bob"));
113    }
114}