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

extendr_api/
iter.rs
