osp_cli/dsl/stages/
limit.rs

1use anyhow::{Result, anyhow};
2
3#[derive(Debug, Clone, Copy, PartialEq, Eq)]
4pub(crate) struct LimitSpec {
5    pub(crate) count: i64,
6    pub(crate) offset: i64,
7}
8
9impl LimitSpec {
10    pub(crate) fn is_head_only(self) -> bool {
11        self.count >= 0 && self.offset >= 0
12    }
13}
14
15pub(crate) fn parse_limit_spec(spec: &str) -> Result<LimitSpec> {
16    let parts: Vec<&str> = spec.split_whitespace().collect();
17    if !(1..=2).contains(&parts.len()) {
18        return Err(anyhow!("L expects 1 or 2 integers (limit [offset])"));
19    }
20
21    let count = parts[0]
22        .parse::<i64>()
23        .map_err(|_| anyhow!("L arguments must be integers"))?;
24    let offset = if parts.len() == 2 {
25        parts[1]
26            .parse::<i64>()
27            .map_err(|_| anyhow!("L arguments must be integers"))?
28    } else {
29        0
30    };
31
32    Ok(LimitSpec { count, offset })
33}
34
35pub fn apply<T>(items: Vec<T>, spec: &str) -> Result<Vec<T>> {
36    let spec = parse_limit_spec(spec)?;
37    let count = spec.count;
38    let offset = spec.offset;
39
40    if count == 0 {
41        return Ok(Vec::new());
42    }
43
44    if count > 0 && offset >= 0 {
45        return Ok(items
46            .into_iter()
47            .skip(offset as usize)
48            .take(count as usize)
49            .collect());
50    }
51
52    let length = items.len() as i64;
53    let start = if offset >= 0 {
54        offset
55    } else {
56        (length + offset).max(0)
57    };
58
59    let base: Vec<T> = items.into_iter().skip(start.max(0) as usize).collect();
60
61    if count >= 0 {
62        Ok(base.into_iter().take(count as usize).collect())
63    } else {
64        let take = count.unsigned_abs() as usize;
65        let skip = base.len().saturating_sub(take);
66        Ok(base.into_iter().skip(skip).collect())
67    }
68}
69
70#[cfg(test)]
71mod tests {
72    use super::apply;
73
74    #[test]
75    fn takes_head_for_positive_limit() {
76        let rows = vec![1, 2, 3];
77        let output = apply(rows, "2").expect("limit should work");
78        assert_eq!(output, vec![1, 2]);
79    }
80
81    #[test]
82    fn handles_zero_limit() {
83        let rows = vec![1, 2, 3];
84        let output = apply(rows, "0").expect("limit should work");
85        assert!(output.is_empty());
86    }
87
88    #[test]
89    fn supports_negative_count_for_tail() {
90        let rows = vec![1, 2, 3, 4, 5];
91        let output = apply(rows, "-2").expect("limit should work");
92        assert_eq!(output, vec![4, 5]);
93    }
94
95    #[test]
96    fn supports_positive_count_with_positive_offset() {
97        let rows = vec![1, 2, 3, 4];
98        let output = apply(rows, "2 1").expect("limit should work");
99        assert_eq!(output, vec![2, 3]);
100    }
101
102    #[test]
103    fn supports_positive_count_with_negative_offset() {
104        let rows = vec![1, 2, 3, 4, 5];
105        let output = apply(rows, "2 -2").expect("limit should work");
106        assert_eq!(output, vec![4, 5]);
107    }
108
109    #[test]
110    fn supports_negative_count_with_negative_offset() {
111        let rows = vec![1, 2, 3, 4, 5];
112        let output = apply(rows, "-1 -2").expect("limit should work");
113        assert_eq!(output, vec![5]);
114    }
115
116    #[test]
117    fn rejects_invalid_argument_count() {
118        let rows = vec![1, 2, 3];
119        let err = apply(rows, "1 2 3").expect_err("invalid arity should fail");
120        assert!(err.to_string().contains("L expects 1 or 2 integers"));
121    }
122}
osp_cli/dsl/stages/limit.rs

osp_cli/dsl/stages/
limit.rs
