extendr_api/
iter.rs

1use crate::*;
2
3use wrapper::symbol::levels_symbol;
4
5/// Iterator over name-value pairs in lists.
6pub type NamedListIter = std::iter::Zip<StrIter, ListIter>;
7
8/// Iterator over strings or string factors.
9///
10/// ```
11/// use extendr_api::prelude::*;
12/// test! {
13///     let robj = r!(["a", "b", "c"]);
14///     assert_eq!(robj.as_str_iter().unwrap().collect::<Vec<_>>(), vec!["a", "b", "c"]);
15///
16///     let factor = factor!(["abcd", "def", "fg", "fg"]);
17///     assert_eq!(factor.levels().unwrap().collect::<Vec<_>>(), vec!["abcd", "def", "fg"]);
18///     assert_eq!(factor.as_integer_vector().unwrap(), vec![1, 2, 3, 3]);
19///     assert_eq!(factor.as_str_iter().unwrap().collect::<Vec<_>>(), vec!["abcd", "def", "fg", "fg"]);
20///     assert_eq!(factor.as_str_iter().unwrap().collect::<Vec<_>>(), vec!["abcd", "def", "fg", "fg"]);
21/// }
22/// ```
23#[derive(Clone)]
24pub struct StrIter {
25    vector: Robj,
26    i: usize,
27    len: usize,
28    levels: SEXP,
29}
30
31impl Default for StrIter {
32    fn default() -> Self {
33        StrIter::new(0)
34    }
35}
36
37impl StrIter {
38    /// Make an empty str iterator.
39    pub fn new(len: usize) -> Self {
40        let vector = if len == 0 { nil_value() } else { na_string() };
41        unsafe {
42            Self {
43                vector,
44                i: 0,
45                len,
46                levels: R_NilValue,
47            }
48        }
49    }
50
51    pub fn na_iter(len: usize) -> StrIter {
52        Self {
53            len,
54            ..Default::default()
55        }
56    }
57}
58
59// Get a string reference from a `CHARSXP`
60pub(crate) fn str_from_strsxp<'a>(sexp: SEXP, index: usize) -> Option<&'a str> {
61    single_threaded(|| unsafe {
62        let charsxp = STRING_ELT(sexp, index as _);
63        rstr::charsxp_to_str(charsxp)
64    })
65}
66
67impl Iterator for StrIter {
68    type Item = &'static str;
69
70    fn size_hint(&self) -> (usize, Option<usize>) {
71        (self.len, Some(self.len))
72    }
73
74    fn next(&mut self) -> Option<Self::Item> {
75        unsafe {
76            let i = self.i;
77            self.i += 1;
78            let vector = self.vector.get();
79            if i >= self.len {
80                None
81            } else if TYPEOF(vector) == SEXPTYPE::NILSXP {
82                None
83            } else if TYPEOF(vector) == SEXPTYPE::STRSXP {
84                str_from_strsxp(vector, i)
85            } else if vector == R_NaString {
86                Some(<&str>::na())
87            } else if TYPEOF(vector) == SEXPTYPE::CHARSXP {
88                rstr::charsxp_to_str(vector)
89            } else if Rf_isFactor(vector).into() {
90                // factor support: factor is an integer, and we need
91                // the value of it, to retrieve the assigned label
92                let level_index = std::slice::from_raw_parts(INTEGER(vector), self.len as _);
93                let level_index = level_index.get(i)?;
94                let level_index = level_index
95                    .checked_sub(1)
96                    .expect("the factor integer has an invalid value in it");
97                str_from_strsxp(self.levels, level_index as _)
98            } else {
99                None
100            }
101        }
102    }
103
104    fn nth(&mut self, n: usize) -> Option<Self::Item> {
105        self.i += n;
106        self.next()
107    }
108}
109
110impl ExactSizeIterator for StrIter {
111    fn len(&self) -> usize {
112        self.len - self.i
113    }
114}
115
116macro_rules! impl_iter_debug {
117    ($name: ty) => {
118        impl std::fmt::Debug for $name {
119            fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
120                write!(f, "[")?;
121                let mut comma = "";
122                for s in self.clone() {
123                    write!(f, "{}{:?}", comma, s)?;
124                    comma = ", ";
125                }
126                write!(f, "]")
127            }
128        }
129    };
130}
131
132impl_iter_debug!(ListIter);
133impl_iter_debug!(PairlistIter);
134impl_iter_debug!(StrIter);
135impl_iter_debug!(EnvIter);
136
137// Lets us create a StrIter from an Robj, e.g. Strings or a factor
138impl TryFrom<&Robj> for StrIter {
139    type Error = Error;
140
141    fn try_from(value: &Robj) -> Result<Self> {
142        value
143            .as_str_iter()
144            .ok_or_else(|| Error::ExpectedString(value.clone()))
145    }
146}
147
148pub trait AsStrIter: GetSexp + Types + Length + Attributes + Rinternals {
149    /// Get an iterator over a string vector.
150    /// Returns None if the object is not a string vector
151    /// but works for factors.
152    ///
153    /// ```
154    /// use extendr_api::prelude::*;
155    ///
156    /// test! {
157    ///     let obj = Robj::from(vec!["a", "b", "c"]);
158    ///     assert_eq!(obj.as_str_iter().unwrap().collect::<Vec<_>>(), vec!["a", "b", "c"]);
159    ///
160    ///     let factor = factor!(vec!["abcd", "def", "fg", "fg"]);
161    ///     assert_eq!(factor.levels().unwrap().collect::<Vec<_>>(), vec!["abcd", "def", "fg"]);
162    ///     assert_eq!(factor.as_integer_vector().unwrap(), vec![1, 2, 3, 3]);
163    ///     assert_eq!(factor.as_str_iter().unwrap().collect::<Vec<_>>(), vec!["abcd", "def", "fg", "fg"]);
164    ///     assert_eq!(factor.as_str_iter().unwrap().collect::<Vec<_>>(), vec!["abcd", "def", "fg", "fg"]);
165    ///
166    ///     let obj = Robj::from(vec![Some("a"), Some("b"), None]);
167    ///     assert_eq!(obj.as_str_iter().unwrap().map(|s| s.is_na()).collect::<Vec<_>>(), vec![false, false, true]);
168    ///
169    ///     let obj = Robj::from(vec!["a", "b", <&str>::na()]);
170    ///     assert_eq!(obj.as_str_iter().unwrap().map(|s| s.is_na()).collect::<Vec<_>>(), vec![false, false, true]);
171    ///
172    ///     let obj = Robj::from(vec!["a", "b", "NA"]);
173    ///     assert_eq!(obj.as_str_iter().unwrap().map(|s| s.is_na()).collect::<Vec<_>>(), vec![false, false, false]);
174    /// }
175    /// ```
176    fn as_str_iter(&self) -> Option<StrIter> {
177        let i = 0;
178        let len = self.len();
179        if self.sexptype() == SEXPTYPE::STRSXP {
180            unsafe {
181                Some(StrIter {
182                    vector: self.as_robj().clone(),
183                    i,
184                    len,
185                    levels: R_NilValue,
186                })
187            }
188        } else if self.sexptype() == SEXPTYPE::CHARSXP {
189            let len = 1;
190            unsafe {
191                Some(StrIter {
192                    vector: self.as_robj().clone(),
193                    i,
194                    len,
195                    levels: R_NilValue,
196                })
197            }
198        } else if self.is_factor() {
199            let levels = self.get_attrib(levels_symbol()).unwrap();
200            unsafe {
201                Some(StrIter {
202                    vector: self.as_robj().clone(),
203                    i,
204                    len,
205                    levels: levels.get(),
206                })
207            }
208        } else {
209            None
210        }
211    }
212}
213
214impl AsStrIter for Robj {}
215
216#[cfg(test)]
217mod tests {
218    use extendr_engine::with_r;
219
220    use super::*;
221
222    #[test]
223    fn single_charsxp_iterator() {
224        with_r(|| {
225            let single_charsxp = blank_string();
226            let s1: Vec<_> = single_charsxp.as_str_iter().unwrap().collect();
227            let single_charsxp = blank_scalar_string();
228            let s2: Vec<_> = single_charsxp.as_str_iter().unwrap().collect();
229            assert_eq!(s1, s2);
230            assert_eq!(s1.len(), 1);
231            assert_eq!(s2.len(), 1);
232        });
233    }
234
235    #[test]
236    fn test_new_constructor() {
237        with_r(|| {
238            let str_iter = StrIter::new(10);
239            assert_eq!(str_iter.collect::<Vec<_>>().len(), 10);
240            let str_iter = StrIter::new(0);
241            let str_iter_collect = str_iter.collect::<Vec<_>>();
242            assert_eq!(str_iter_collect.len(), 0);
243            assert!(str_iter_collect.is_empty());
244            let mut str_iter = StrIter::new(0);
245            assert!(str_iter.next().is_none());
246        });
247    }
248}
extendr_api/iter.rs

extendr_api/
iter.rs
