extendr_api/
iter.rs

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 || TYPEOF(vector) == SEXPTYPE::NILSXP {
80                None
81            } else if TYPEOF(vector) == SEXPTYPE::STRSXP {
82                str_from_strsxp(vector, i)
83            } else if vector == R_NaString {
84                Some(<&str>::na())
85            } else if TYPEOF(vector) == SEXPTYPE::CHARSXP {
86                rstr::charsxp_to_str(vector)
87            } else if Rf_isFactor(vector).into() {
88                // factor support: factor is an integer, and we need
89                // the value of it, to retrieve the assigned label
90                let level_index = std::slice::from_raw_parts(INTEGER(vector), self.len as _);
91                let level_index = level_index.get(i)?;
92                let level_index = level_index
93                    .checked_sub(1)
94                    .expect("the factor integer has an invalid value in it");
95                str_from_strsxp(self.levels, level_index as _)
96            } else {
97                None
98            }
99        }
100    }
101
102    fn nth(&mut self, n: usize) -> Option<Self::Item> {
103        self.i += n;
104        self.next()
105    }
106}
107
108impl ExactSizeIterator for StrIter {
109    fn len(&self) -> usize {
110        self.len - self.i
111    }
112}
113
114macro_rules! impl_iter_debug {
115    ($name: ty) => {
116        impl std::fmt::Debug for $name {
117            fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
118                write!(f, "[")?;
119                let mut comma = "";
120                for s in self.clone() {
121                    write!(f, "{}{:?}", comma, s)?;
122                    comma = ", ";
123                }
124                write!(f, "]")
125            }
126        }
127    };
128}
129
130impl_iter_debug!(ListIter);
131impl_iter_debug!(PairlistIter);
132impl_iter_debug!(StrIter);
133impl_iter_debug!(EnvIter);
134
135// Lets us create a StrIter from an Robj, e.g. Strings or a factor
136impl TryFrom<&Robj> for StrIter {
137    type Error = Error;
138
139    fn try_from(value: &Robj) -> Result<Self> {
140        value
141            .as_str_iter()
142            .ok_or_else(|| Error::ExpectedString(value.clone()))
143    }
144}
145
146impl TryFrom<Robj> for StrIter {
147    type Error = Error;
148
149    fn try_from(value: Robj) -> Result<Self> {
150        (&value).try_into()
151    }
152}
153
154pub trait AsStrIter: GetSexp + Types + Length + Attributes + Rinternals {
155    /// Get an iterator over a string vector.
156    /// Returns None if the object is not a string vector
157    /// but works for factors.
158    ///
159    /// ```
160    /// use extendr_api::prelude::*;
161    ///
162    /// test! {
163    ///     let obj = Robj::from(vec!["a", "b", "c"]);
164    ///     assert_eq!(obj.as_str_iter().unwrap().collect::<Vec<_>>(), vec!["a", "b", "c"]);
165    ///
166    ///     let factor = factor!(vec!["abcd", "def", "fg", "fg"]);
167    ///     assert_eq!(factor.levels().unwrap().collect::<Vec<_>>(), vec!["abcd", "def", "fg"]);
168    ///     assert_eq!(factor.as_integer_vector().unwrap(), vec![1, 2, 3, 3]);
169    ///     assert_eq!(factor.as_str_iter().unwrap().collect::<Vec<_>>(), vec!["abcd", "def", "fg", "fg"]);
170    ///     assert_eq!(factor.as_str_iter().unwrap().collect::<Vec<_>>(), vec!["abcd", "def", "fg", "fg"]);
171    ///
172    ///     let obj = Robj::from(vec![Some("a"), Some("b"), None]);
173    ///     assert_eq!(obj.as_str_iter().unwrap().map(|s| s.is_na()).collect::<Vec<_>>(), vec![false, false, true]);
174    ///
175    ///     let obj = Robj::from(vec!["a", "b", <&str>::na()]);
176    ///     assert_eq!(obj.as_str_iter().unwrap().map(|s| s.is_na()).collect::<Vec<_>>(), vec![false, false, true]);
177    ///
178    ///     let obj = Robj::from(vec!["a", "b", "NA"]);
179    ///     assert_eq!(obj.as_str_iter().unwrap().map(|s| s.is_na()).collect::<Vec<_>>(), vec![false, false, false]);
180    /// }
181    /// ```
182    fn as_str_iter(&self) -> Option<StrIter> {
183        let i = 0;
184        let len = self.len();
185        if self.sexptype() == SEXPTYPE::STRSXP {
186            unsafe {
187                Some(StrIter {
188                    vector: self.as_robj().clone(),
189                    i,
190                    len,
191                    levels: R_NilValue,
192                })
193            }
194        } else if self.sexptype() == SEXPTYPE::CHARSXP {
195            let len = 1;
196            unsafe {
197                Some(StrIter {
198                    vector: self.as_robj().clone(),
199                    i,
200                    len,
201                    levels: R_NilValue,
202                })
203            }
204        } else if self.is_factor() {
205            let levels = self.get_attrib(levels_symbol()).unwrap();
206            unsafe {
207                Some(StrIter {
208                    vector: self.as_robj().clone(),
209                    i,
210                    len,
211                    levels: levels.get(),
212                })
213            }
214        } else {
215            None
216        }
217    }
218}
219
220impl AsStrIter for Robj {}
221
222#[cfg(test)]
223mod tests {
224    use extendr_engine::with_r;
225
226    use super::*;
227
228    #[test]
229    fn single_charsxp_iterator() {
230        with_r(|| {
231            let single_charsxp = blank_string();
232            let s1: Vec<_> = single_charsxp.as_str_iter().unwrap().collect();
233            let single_charsxp = blank_scalar_string();
234            let s2: Vec<_> = single_charsxp.as_str_iter().unwrap().collect();
235            assert_eq!(s1, s2);
236            assert_eq!(s1.len(), 1);
237            assert_eq!(s2.len(), 1);
238        });
239    }
240
241    #[test]
242    fn test_new_constructor() {
243        with_r(|| {
244            let str_iter = StrIter::new(10);
245            assert_eq!(str_iter.collect::<Vec<_>>().len(), 10);
246            let str_iter = StrIter::new(0);
247            let str_iter_collect = str_iter.collect::<Vec<_>>();
248            assert_eq!(str_iter_collect.len(), 0);
249            assert!(str_iter_collect.is_empty());
250            let mut str_iter = StrIter::new(0);
251            assert!(str_iter.next().is_none());
252        });
253    }
254}
extendr_api/iter.rs

extendr_api/
iter.rs
