Skip to content

Commit e2d9b11

Browse files
mitsuhikomitsuhiko
authored
Implement negative slices and fixed byte slices (#781)
1 parent f50a7c2 commit e2d9b11

File tree

1 file changed

+238
-23
lines changed

1 file changed

+238
-23
lines changed

‎minijinja/src/value/ops.rs

Lines changed: 238 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -87,8 +87,41 @@ fn get_offset_and_len<F: FnOnce() -> usize>(
8787
}
8888
}
8989

90+
fn range_step_backwards(
91+
start_was_none: bool,
92+
start: i64,
93+
stop: Option<i64>,
94+
step: usize,
95+
end: usize,
96+
) -> impl Iterator<Item = usize> {
97+
let start = if start_was_none {
98+
end.saturating_sub(1)
99+
} else if start >= 0 {
100+
if start as usize >= end {
101+
end.saturating_sub(1)
102+
} else {
103+
start as usize
104+
}
105+
} else {
106+
(end as i64 + start).max(0) as usize
107+
};
108+
let stop = match stop {
109+
None => 0,
110+
Some(stop) if stop < 0 => (end as i64 + stop).max(0) as usize,
111+
Some(stop) => stop as usize,
112+
};
113+
let length = if stop == 0 {
114+
(start + step) / step
115+
} else {
116+
(start - stop + step - 1) / step
117+
};
118+
std::iter::successors(Some(start), move |&i| i.checked_sub(step)).take(length)
119+
}
120+
90121
pub fn slice(value: Value, start: Value, stop: Value, step: Value) -> Result<Value, Error> {
91-
let start: i64 = if start.is_none() {
122+
let start_was_none = start.is_none();
123+
124+
let start: i64 = if start_was_none {
92125
0
93126
} else {
94127
ok!(start.try_into())
@@ -99,9 +132,9 @@ pub fn slice(value: Value, start: Value, stop: Value, step: Value) -> Result<Val
99132
Some(ok!(i64::try_from(stop)))
100133
};
101134
let step = if step.is_none() {
102-
1
135+
1i64
103136
} else {
104-
ok!(u64::try_from(step)) as usize
137+
ok!(i64::try_from(step))
105138
};
106139
if step == 0 {
107140
return Err(Error::new(
@@ -119,33 +152,74 @@ pub fn slice(value: Value, start: Value, stop: Value, step: Value) -> Result<Val
119152
match value.0 {
120153
ValueRepr::String(..) | ValueRepr::SmallStr(_) => {
121154
let s = value.as_str().unwrap();
122-
let (start, len) = get_offset_and_len(start, stop, || s.chars().count());
123-
Ok(Value::from(
124-
s.chars()
125-
.skip(start)
126-
.take(len)
127-
.step_by(step)
128-
.collect::<String>(),
129-
))
155+
if step > 0 {
156+
let (start, len) = get_offset_and_len(start, stop, || s.chars().count());
157+
Ok(Value::from(
158+
s.chars()
159+
.skip(start)
160+
.take(len)
161+
.step_by(step as usize)
162+
.collect::<String>(),
163+
))
164+
} else {
165+
let chars: Vec<char> = s.chars().collect();
166+
Ok(Value::from(
167+
range_step_backwards(start_was_none, start, stop, -step as usize, chars.len())
168+
.map(move |i| chars[i])
169+
.collect::<String>(),
170+
))
171+
}
130172
}
131173
ValueRepr::Bytes(ref b) => {
132-
let (start, len) = get_offset_and_len(start, stop, || b.len());
133-
Ok(Value::from_bytes(
134-
b.get(start..start + len).unwrap_or_default().to_owned(),
135-
))
174+
if step > 0 {
175+
let (start, len) = get_offset_and_len(start, stop, || b.len());
176+
Ok(Value::from_bytes(
177+
b.iter()
178+
.skip(start)
179+
.take(len)
180+
.step_by(step as usize)
181+
.copied()
182+
.collect(),
183+
))
184+
} else {
185+
Ok(Value::from_bytes(
186+
range_step_backwards(start_was_none, start, stop, -step as usize, b.len())
187+
.map(|i| b[i])
188+
.collect::<Vec<u8>>(),
189+
))
190+
}
136191
}
137192
ValueRepr::Undefined(_) | ValueRepr::None => Ok(Value::from(Vec::<Value>::new())),
138193
ValueRepr::Object(obj) if matches!(obj.repr(), ObjectRepr::Seq | ObjectRepr::Iterable) => {
139-
Ok(Value::make_object_iterable(obj, move |obj| {
194+
if step > 0 {
140195
let len = obj.enumerator_len().unwrap_or_default();
141196
let (start, len) = get_offset_and_len(start, stop, || len);
142-
// The manual match here is important that we do not mess up the size_hint
143-
if let Some(iter) = obj.try_iter() {
144-
Box::new(iter.skip(start).take(len).step_by(step))
145-
} else {
146-
Box::new(None.into_iter())
147-
}
148-
}))
197+
Ok(Value::make_object_iterable(obj, move |obj| {
198+
if let Some(iter) = obj.try_iter() {
199+
Box::new(iter.skip(start).take(len).step_by(step as usize))
200+
} else {
201+
Box::new(None.into_iter())
202+
}
203+
}))
204+
} else {
205+
Ok(Value::make_object_iterable(obj.clone(), move |obj| {
206+
if let Some(iter) = obj.try_iter() {
207+
let vec: Vec<Value> = iter.collect();
208+
Box::new(
209+
range_step_backwards(
210+
start_was_none,
211+
start,
212+
stop,
213+
-step as usize,
214+
vec.len(),
215+
)
216+
.map(move |i| vec[i].clone()),
217+
)
218+
} else {
219+
Box::new(None.into_iter())
220+
}
221+
}))
222+
}
149223
}
150224
_ => error,
151225
}
@@ -500,4 +574,145 @@ mod tests {
500574
Value::from("2342")
501575
);
502576
}
577+
578+
#[test]
579+
fn test_slicing() {
580+
let v = Value::from(vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9]);
581+
582+
// [::] - full slice
583+
assert_eq!(
584+
slice(v.clone(), Value::from(()), Value::from(()), Value::from(())).unwrap(),
585+
Value::from(vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
586+
);
587+
588+
// [::2] - every 2nd element
589+
assert_eq!(
590+
slice(v.clone(), Value::from(()), Value::from(()), Value::from(2)).unwrap(),
591+
Value::from(vec![0, 2, 4, 6, 8])
592+
);
593+
594+
// [1:2:2] - slice with start, stop, step
595+
assert_eq!(
596+
slice(v.clone(), Value::from(1), Value::from(2), Value::from(2)).unwrap(),
597+
Value::from(vec![1])
598+
);
599+
600+
// [::-2] - reverse with step of 2
601+
assert_eq!(
602+
slice(v.clone(), Value::from(()), Value::from(()), Value::from(-2)).unwrap(),
603+
Value::from(vec![9, 7, 5, 3, 1])
604+
);
605+
606+
// [2::-2] - from index 2 to start, reverse with step of 2
607+
assert_eq!(
608+
slice(v.clone(), Value::from(2), Value::from(()), Value::from(-2)).unwrap(),
609+
Value::from(vec![2, 0])
610+
);
611+
612+
// [4:2:-2] - from index 4 to 2, reverse with step of 2
613+
assert_eq!(
614+
slice(v.clone(), Value::from(4), Value::from(2), Value::from(-2)).unwrap(),
615+
Value::from(vec![4])
616+
);
617+
618+
// [8:3:-2] - from index 8 to 3, reverse with step of 2
619+
assert_eq!(
620+
slice(v.clone(), Value::from(8), Value::from(3), Value::from(-2)).unwrap(),
621+
Value::from(vec![8, 6, 4])
622+
);
623+
}
624+
625+
#[test]
626+
fn test_string_slicing() {
627+
let s = Value::from("abcdefghij");
628+
629+
// [::] - full slice
630+
assert_eq!(
631+
slice(s.clone(), Value::from(()), Value::from(()), Value::from(())).unwrap(),
632+
Value::from("abcdefghij")
633+
);
634+
635+
// [::2] - every 2nd character
636+
assert_eq!(
637+
slice(s.clone(), Value::from(()), Value::from(()), Value::from(2)).unwrap(),
638+
Value::from("acegi")
639+
);
640+
641+
// [1:2:2] - slice with start, stop, step
642+
assert_eq!(
643+
slice(s.clone(), Value::from(1), Value::from(2), Value::from(2)).unwrap(),
644+
Value::from("b")
645+
);
646+
647+
// [::-2] - reverse with step of 2
648+
assert_eq!(
649+
slice(s.clone(), Value::from(()), Value::from(()), Value::from(-2)).unwrap(),
650+
Value::from("jhfdb")
651+
);
652+
653+
// [2::-2] - from index 2 to start, reverse with step of 2
654+
assert_eq!(
655+
slice(s.clone(), Value::from(2), Value::from(()), Value::from(-2)).unwrap(),
656+
Value::from("ca")
657+
);
658+
659+
// [4:2:-2] - from index 4 to 2, reverse with step of 2
660+
assert_eq!(
661+
slice(s.clone(), Value::from(4), Value::from(2), Value::from(-2)).unwrap(),
662+
Value::from("e")
663+
);
664+
665+
// [8:3:-2] - from index 8 to 3, reverse with step of 2
666+
assert_eq!(
667+
slice(s.clone(), Value::from(8), Value::from(3), Value::from(-2)).unwrap(),
668+
Value::from("ige")
669+
);
670+
}
671+
672+
#[test]
673+
fn test_bytes_slicing() {
674+
let s = Value::from_bytes(b"abcdefghij".to_vec());
675+
676+
// [::] - full slice
677+
assert_eq!(
678+
slice(s.clone(), Value::from(()), Value::from(()), Value::from(())).unwrap(),
679+
Value::from_bytes(b"abcdefghij".to_vec())
680+
);
681+
682+
// [::2] - every 2nd character
683+
assert_eq!(
684+
slice(s.clone(), Value::from(()), Value::from(()), Value::from(2)).unwrap(),
685+
Value::from_bytes(b"acegi".to_vec())
686+
);
687+
688+
// [1:2:2] - slice with start, stop, step
689+
assert_eq!(
690+
slice(s.clone(), Value::from(1), Value::from(2), Value::from(2)).unwrap(),
691+
Value::from_bytes(b"b".to_vec())
692+
);
693+
694+
// [::-2] - reverse with step of 2
695+
assert_eq!(
696+
slice(s.clone(), Value::from(()), Value::from(()), Value::from(-2)).unwrap(),
697+
Value::from_bytes(b"jhfdb".to_vec())
698+
);
699+
700+
// [2::-2] - from index 2 to start, reverse with step of 2
701+
assert_eq!(
702+
slice(s.clone(), Value::from(2), Value::from(()), Value::from(-2)).unwrap(),
703+
Value::from_bytes(b"ca".to_vec())
704+
);
705+
706+
// [4:2:-2] - from index 4 to 2, reverse with step of 2
707+
assert_eq!(
708+
slice(s.clone(), Value::from(4), Value::from(2), Value::from(-2)).unwrap(),
709+
Value::from_bytes(b"e".to_vec())
710+
);
711+
712+
// [8:3:-2] - from index 8 to 3, reverse with step of 2
713+
assert_eq!(
714+
slice(s.clone(), Value::from(8), Value::from(3), Value::from(-2)).unwrap(),
715+
Value::from_bytes(b"ige".to_vec())
716+
);
717+
}
503718
}

0 commit comments

Comments
 (0)