Added chain filter (#807)

Author: mitsuhiko

CHANGELOG.md

@@ -13,6 +13,7 @@ All notable changes to MiniJinja are documented here.
 * Allow `isize` as argument type.  #799
 * MiniJinja now correctly handles `\x` escape sequences in strings
   as well as octals.  #805
+* Added a new `|chain` filter.  #807
 
 ## 2.10.2
 

minijinja/src/defaults.rs

@@ -110,6 +110,7 @@ pub(crate) fn get_builtin_filters() -> BTreeMap<Cow<'static, str>, Value> {
         rv.insert("map".into(), Value::from_function(filters::map));
         rv.insert("groupby".into(), Value::from_function(filters::groupby));
         rv.insert("unique".into(), Value::from_function(filters::unique));
+        rv.insert("chain".into(), Value::from_function(filters::chain));
         rv.insert("pprint".into(), Value::from_function(filters::pprint));
 
         #[cfg(feature = "json")]

minijinja/src/filters.rs

@@ -113,8 +113,6 @@
 //!
 //! Some additional filters are available in the
 //! [`minijinja-contrib`](https://crates.io/crates/minijinja-contrib) crate.
-use std::sync::Arc;
-
 use crate::error::Error;
 use crate::utils::write_escaped;
 use crate::value::Value;
@@ -169,12 +167,14 @@ mod builtins {
 
     use crate::error::ErrorKind;
     use crate::utils::{safe_sort, splitn_whitespace};
+    use crate::value::merge_object::{MergeDict, MergeSeq};
     use crate::value::ops::{self, as_f64};
     use crate::value::{Enumerator, Kwargs, Object, ObjectRepr, ValueKind, ValueRepr};
     use std::borrow::Cow;
     use std::cmp::Ordering;
     use std::fmt::Write;
     use std::mem;
+    use std::sync::Arc;
 
     /// Converts a value to uppercase.
     ///
@@ -1467,6 +1467,57 @@ mod builtins {
         Ok(Value::from(rv))
     }
 
+    /// Chain two or more iterable objects as a single iterable object.
+    ///
+    /// If all the individual objects are dictionaries, then the final chained object
+    /// also acts like a dictionary -- you can lookup a key, or iterate over the keys
+    /// etc. Note that the dictionaries are not merged, so if there are duplicate keys,
+    /// then the lookup will return the value from the last matching dictionary in the
+    /// chain.
+    ///
+    /// If all the individual objects are sequences, then the final chained
+    /// object also acts like a list as if the lists are appended.
+    ///
+    /// Otherwise, the chained object acts like an iterator chaining individual
+    /// iterators, but it cannot be indexed.
+    ///
+    /// ```jinja
+    /// {{ users | chain(moreusers) | length }}
+    /// {% for user, info in shard0 | chain(shard1, shard2) | dictsort %}
+    ///   {{user}}: {{info}}
+    /// {% endfor %}
+    /// {{ list1 | chain(list2) | attr(1) }}
+    /// ```
+    #[cfg_attr(docsrs, doc(cfg(feature = "builtins")))]
+    pub fn chain(
+        _state: &State,
+        value: Value,
+        others: crate::value::Rest<Value>,
+    ) -> Result<Value, Error> {
+        let all_values = Some(value.clone())
+            .into_iter()
+            .chain(others.0.iter().cloned())
+            .collect::<Vec<_>>();
+
+        if all_values.iter().all(|v| v.kind() == ValueKind::Map) {
+            Ok(Value::from_object(MergeDict::new(all_values)))
+        } else if all_values
+            .iter()
+            .all(|v| matches!(v.kind(), ValueKind::Seq))
+        {
+            Ok(Value::from_object(MergeSeq::new(all_values)))
+        } else {
+            // General iterator chaining behavior
+            Ok(Value::make_object_iterable(all_values, |values| {
+                Box::new(values.iter().flat_map(|v| match v.try_iter() {
+                    Ok(iter) => Box::new(iter) as Box<dyn Iterator<Item = Value> + Send + Sync>,
+                    Err(err) => Box::new(Some(Value::from(err)).into_iter())
+                        as Box<dyn Iterator<Item = Value> + Send + Sync>,
+                })) as Box<dyn Iterator<Item = Value> + Send + Sync>
+            }))
+        }
+    }
+
     /// Pretty print a variable.
     ///
     /// This is useful for debugging as it better shows what's inside an object.

minijinja/src/value/merge_object.rs

@@ -1,30 +1,97 @@
 use std::collections::BTreeSet;
 use std::sync::Arc;
 
-use crate::value::{Enumerator, Object, Value, ValueKind};
+use crate::value::ops::LenIterWrap;
+use crate::value::{Enumerator, Object, ObjectExt, ObjectRepr, Value, ValueKind};
 
-#[derive(Clone, Debug)]
-struct MergeObject(pub Box<[Value]>);
+/// Dictionary merging behavior - create custom object with lookup capability
+#[derive(Debug)]
+pub struct MergeDict {
+    values: Box<[Value]>,
+}
 
-impl Object for MergeObject {
+impl MergeDict {
+    pub fn new(values: Vec<Value>) -> Self {
+        Self {
+            values: values.into_boxed_slice(),
+        }
+    }
+}
+
+impl Object for MergeDict {
     fn get_value(self: &Arc<Self>, key: &Value) -> Option<Value> {
-        self.0
-            .iter()
-            .filter_map(|x| x.get_item_opt(key))
-            .find(|x| !x.is_undefined())
+        // Look up key in reverse order (last matching dict wins)
+        for value in self.values.iter().rev() {
+            if let Ok(v) = value.get_item(key) {
+                if !v.is_undefined() {
+                    return Some(v);
+                }
+            }
+        }
+        None
     }
 
     fn enumerate(self: &Arc<Self>) -> Enumerator {
-        // we collect here the whole internal object once on iteration so that
-        // we have an enumerator with a known length.
-        let items = self
-            .0
+        // Collect all keys from all dictionaries (only include maps)
+        let keys: BTreeSet<Value> = self
+            .values
             .iter()
             .filter(|x| x.kind() == ValueKind::Map)
-            .flat_map(|v| v.try_iter().ok())
+            .filter_map(|v| v.try_iter().ok())
             .flatten()
-            .collect::<BTreeSet<_>>();
-        Enumerator::Iter(Box::new(items.into_iter()))
+            .collect();
+        Enumerator::Iter(Box::new(keys.into_iter()))
+    }
+}
+
+/// List merging behavior - calculate total length for size hint
+#[derive(Debug)]
+pub struct MergeSeq {
+    values: Box<[Value]>,
+    total_len: Option<usize>,
+}
+
+impl MergeSeq {
+    pub fn new(values: Vec<Value>) -> Self {
+        Self {
+            total_len: values.iter().map(|v| v.len()).sum(),
+            values: values.into_boxed_slice(),
+        }
+    }
+}
+
+impl Object for MergeSeq {
+    fn repr(self: &Arc<Self>) -> ObjectRepr {
+        ObjectRepr::Seq
+    }
+
+    fn get_value(self: &Arc<Self>, key: &Value) -> Option<Value> {
+        if let Some(idx) = key.as_usize() {
+            let mut current_idx = 0;
+            for value in self.values.iter() {
+                let len = value.len().unwrap_or(0);
+                if idx < current_idx + len {
+                    return value.get_item(&Value::from(idx - current_idx)).ok();
+                }
+                current_idx += len;
+            }
+        }
+        None
+    }
+
+    fn enumerate(self: &Arc<Self>) -> Enumerator {
+        self.mapped_enumerator(|this| {
+            let iter = this.values.iter().flat_map(|v| match v.try_iter() {
+                Ok(iter) => Box::new(iter) as Box<dyn Iterator<Item = Value> + Send + Sync>,
+                Err(err) => Box::new(Some(Value::from(err)).into_iter())
+                    as Box<dyn Iterator<Item = Value> + Send + Sync>,
+            });
+            if let Some(total_len) = this.total_len {
+                Box::new(LenIterWrap(total_len, iter))
+            } else {
+                Box::new(iter)
+            }
+        })
     }
 }
 
@@ -60,14 +127,13 @@ where
     I: IntoIterator<Item = V>,
     V: Into<Value>,
 {
-    let mut sources: Box<[Value]> = iter.into_iter().map(Into::into).collect();
+    let sources: Vec<Value> = iter.into_iter().map(Into::into).collect();
     // if we only have a single source, we can use it directly to avoid making
     // an unnecessary indirection.
     if sources.len() == 1 {
         sources[0].clone()
     } else {
-        sources.reverse();
-        Value::from_object(MergeObject(sources))
+        Value::from_object(MergeDict::new(sources))
     }
 }
 

minijinja/src/value/ops.rs

@@ -3,6 +3,23 @@ use crate::value::{DynObject, ObjectRepr, Value, ValueKind, ValueRepr};
 
 const MIN_I128_AS_POS_U128: u128 = 170141183460469231731687303715884105728;
 
+/// Iterator wrapper that provides exact size hints for iterators with known length.
+pub(crate) struct LenIterWrap<I: Send + Sync>(pub(crate) usize, pub(crate) I);
+
+impl<I: Iterator<Item = Value> + Send + Sync> Iterator for LenIterWrap<I> {
+    type Item = Value;
+
+    #[inline(always)]
+    fn next(&mut self) -> Option<Self::Item> {
+        self.1.next()
+    }
+
+    #[inline(always)]
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        (self.0, Some(self.0))
+    }
+}
+
 pub enum CoerceResult<'a> {
     I128(i128, i128),
     F64(f64, f64),
@@ -316,22 +333,6 @@ pub fn mul(lhs: &Value, rhs: &Value) -> Result<Value, Error> {
 }
 
 fn repeat_iterable(n: &Value, seq: &DynObject) -> Result<Value, Error> {
-    struct LenIterWrap<I: Send + Sync>(usize, I);
-
-    impl<I: Iterator<Item = Value> + Send + Sync> Iterator for LenIterWrap<I> {
-        type Item = Value;
-
-        #[inline(always)]
-        fn next(&mut self) -> Option<Self::Item> {
-            self.1.next()
-        }
-
-        #[inline(always)]
-        fn size_hint(&self) -> (usize, Option<usize>) {
-            (self.0, Some(self.0))
-        }
-    }
-
     let n = ok!(n.as_usize().ok_or_else(|| {
         Error::new(
             ErrorKind::InvalidOperation,

minijinja/tests/inputs/filters.txt

@@ -16,7 +16,13 @@
     "a": "b",
     "c": "d"
   },
-  "scary_html": "<>&'"
+  "scary_html": "<>&'",
+  "list2": [4, 5, 6],
+  "list3": [7, 8],
+  "map2": {
+    "e": "f",
+    "g": "h"
+  }
 }
 ---
 lower: {{ word|lower }}
@@ -121,3 +127,9 @@ sum-empty: {{ []|sum }}
 sum-float: {{ [0.5, 1.0]|sum }}
 lines: {{ "foo\nbar\r\nbaz"|lines }}
 string: {{ [1|string, 2|string] }}
+chain-lists: {{ list|chain(list2)|list }}
+chain-multiple: {{ list|chain(list2, list3)|list }}
+chain-maps: {{ map|chain(map2) }}
+chain-and-length: {{ list|chain(list2)|length }}
+chain-iteration: {% for item in list|chain(list2) %}{{ item }}{% if not loop.last %},{% endif %}{% endfor %}
+chain-indexing: {{ (list|chain(list2))[4] }}

minijinja/tests/snapshots/[email protected]

@@ -74,6 +74,7 @@ State {
             "batch",
             "bool",
             "capitalize",
+            "chain",
             "count",
             "d",
             "default",