extendr_api/
iter.rs

Help
1use crate::*;
2
3use extendr_ffi::{R_NaString, R_NilValue, Rf_isFactor, INTEGER, STRING_ELT, TYPEOF};
4use wrapper::symbol::levels_symbol;
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
148impl TryFrom<Robj> for StrIter {
149    type Error = Error;
150
151    fn try_from(value: Robj) -> Result<Self> {
152        (&value).try_into()
153    }
154}
155
156pub trait AsStrIter: GetSexp + Types + Length + Attributes + Rinternals {
157    /// Get an iterator over a string vector.
158    /// Returns None if the object is not a string vector
159    /// but works for factors.
160    ///
161    /// ```
162    /// use extendr_api::prelude::*;
163    ///
164    /// test! {
165    ///     let obj = Robj::from(vec!["a", "b", "c"]);
166    ///     assert_eq!(obj.as_str_iter().unwrap().collect::<Vec<_>>(), vec!["a", "b", "c"]);
167    ///
168    ///     let factor = factor!(vec!["abcd", "def", "fg", "fg"]);
169    ///     assert_eq!(factor.levels().unwrap().collect::<Vec<_>>(), vec!["abcd", "def", "fg"]);
170    ///     assert_eq!(factor.as_integer_vector().unwrap(), vec![1, 2, 3, 3]);
171    ///     assert_eq!(factor.as_str_iter().unwrap().collect::<Vec<_>>(), vec!["abcd", "def", "fg", "fg"]);
172    ///     assert_eq!(factor.as_str_iter().unwrap().collect::<Vec<_>>(), vec!["abcd", "def", "fg", "fg"]);
173    ///
174    ///     let obj = Robj::from(vec![Some("a"), Some("b"), None]);
175    ///     assert_eq!(obj.as_str_iter().unwrap().map(|s| s.is_na()).collect::<Vec<_>>(), vec![false, false, true]);
176    ///
177    ///     let obj = Robj::from(vec!["a", "b", <&str>::na()]);
178    ///     assert_eq!(obj.as_str_iter().unwrap().map(|s| s.is_na()).collect::<Vec<_>>(), vec![false, false, true]);
179    ///
180    ///     let obj = Robj::from(vec!["a", "b", "NA"]);
181    ///     assert_eq!(obj.as_str_iter().unwrap().map(|s| s.is_na()).collect::<Vec<_>>(), vec![false, false, false]);
182    /// }
183    /// ```
184    fn as_str_iter(&self) -> Option<StrIter> {
185        let i = 0;
186        let len = self.len();
187        if self.sexptype() == SEXPTYPE::STRSXP {
188            unsafe {
189                Some(StrIter {
190                    vector: self.as_robj().clone(),
191                    i,
192                    len,
193                    levels: R_NilValue,
194                })
195            }
196        } else if self.sexptype() == SEXPTYPE::CHARSXP {
197            let len = 1;
198            unsafe {
199                Some(StrIter {
200                    vector: self.as_robj().clone(),
201                    i,
202                    len,
203                    levels: R_NilValue,
204                })
205            }
206        } else if self.is_factor() {
207            let levels = self.get_attrib(levels_symbol()).unwrap();
208            unsafe {
209                Some(StrIter {
210                    vector: self.as_robj().clone(),
211                    i,
212                    len,
213                    levels: levels.get(),
214                })
215            }
216        } else {
217            None
218        }
219    }
220}
221
222impl AsStrIter for Robj {}
223
224#[cfg(test)]
225mod tests {
226    use extendr_engine::with_r;
227
228    use super::*;
229
230    #[test]
231    fn single_charsxp_iterator() {
232        with_r(|| {
233            let single_charsxp = blank_string();
234            let s1: Vec<_> = single_charsxp.as_str_iter().unwrap().collect();
235            let single_charsxp = blank_scalar_string();
236            let s2: Vec<_> = single_charsxp.as_str_iter().unwrap().collect();
237            assert_eq!(s1, s2);
238            assert_eq!(s1.len(), 1);
239            assert_eq!(s2.len(), 1);
240        });
241    }
242
243    #[test]
244    fn test_new_constructor() {
245        with_r(|| {
246            let str_iter = StrIter::new(10);
247            assert_eq!(str_iter.collect::<Vec<_>>().len(), 10);
248            let str_iter = StrIter::new(0);
249            let str_iter_collect = str_iter.collect::<Vec<_>>();
250            assert_eq!(str_iter_collect.len(), 0);
251            assert!(str_iter_collect.is_empty());
252            let mut str_iter = StrIter::new(0);
253            assert!(str_iter.next().is_none());
254        });
255    }
256}
extendr_api/iter.rs

extendr_api/
iter.rs
