Repository: mitsuhiko/minijinja
Pull request #521 (2024-06-09 22:42:02)

Adds python compatibility support to contrib
This adds `pycompat`, a new module to `minijinja-contrib`. It also adds a `--py-compat` flag to the CLI which enables it.

Once this is registered as an unknown method callback the engine supports functions such as `dict.keys` for better compatibility with Python based Jinja2 templates.

Refs #438


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 4db3a8a..e62cfa6 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -6,6 +6,11 @@ All notable changes to MiniJinja are documented here.
 
 - Implemented sequence (+ some iterator) and string repeating with the `*`
   operator to match Jinja2 behavior.  #519
+- Added the new `minijinja::pycompat` module which allows one to register
+  an unknown method callback that provides most built-in Python methods.
+  This makes things such as `dict.keys` work.  Also adds a new
+  `--py-compat` flag to `minijinja-cli` that enables it.  This improves
+  the compatibility with Python based templates.  #521
 
 ## 2.0.1
 
diff --git a/COMPATIBILITY.md b/COMPATIBILITY.md
index 971ea3e..a915a0d 100644
--- a/COMPATIBILITY.md
+++ b/COMPATIBILITY.md
@@ -27,6 +27,10 @@ you cannot do `x.items()` to iterate over items.  For this particular case
 both Jinja2 and MiniJinja now support `|items` for iteration instead.  Other
 methods are rarely useful and filters should be used instead.
 
+Support for these Python methods can however be loaded by registering the
+`unknown_method_callback` from the `pycompat` module in the `minijinja-contrib`
+crate.
+
 ### Tuples
 
 MiniJinja does not implement tuples.  The creation of tuples with tuple syntax
diff --git a/minijinja-cli/Cargo.toml b/minijinja-cli/Cargo.toml
index c455a37..fdd8acf 100644
--- a/minijinja-cli/Cargo.toml
+++ b/minijinja-cli/Cargo.toml
@@ -41,7 +41,7 @@ minijinja = { version = "=2.0.1", path = "../minijinja", features = [
     "unstable_machinery",
     "custom_syntax",
 ] }
-minijinja-contrib = { version = "=2.0.1", path = "../minijinja-contrib" }
+minijinja-contrib = { version = "=2.0.1", path = "../minijinja-contrib", features = ["pycompat"] }
 rustyline = { version = "12.0.0", optional = true }
 serde = "1.0.183"
 serde_json = "1.0.105"
diff --git a/minijinja-cli/README.md b/minijinja-cli/README.md
index c38332b..4721464 100644
--- a/minijinja-cli/README.md
+++ b/minijinja-cli/README.md
@@ -77,6 +77,9 @@ can be set to stdin at once.
     Enable the trim_blocks flag
 - `--lstrip-blocks`:
     Enable the lstrip_blocks flag
+- `--py-compat`:
+    Enables improved Python compatibility.  Enabling this adds methods such as
+    `dict.keys` and some others.
 - `-s`, `--syntax <PAIR>`:
     Changes a syntax feature (feature=value) [possible features: `block-start`, `block-end`, `variable-start`, `variable-end`, `comment-start`, `comment-end`, `line-statement-prefix`, `line-statement-comment`]
 - `--safe-path <PATH>`:
diff --git a/minijinja-cli/src/cli.rs b/minijinja-cli/src/cli.rs
index 59d521a..d413db5 100644
--- a/minijinja-cli/src/cli.rs
+++ b/minijinja-cli/src/cli.rs
@@ -29,6 +29,8 @@ pub(super) fn make_command() -> Command {
             arg!(--"no-newline" "Do not output a trailing newline"),
             arg!(--"trim-blocks" "Enable the trim_blocks flag"),
             arg!(--"lstrip-blocks" "Enable the lstrip_blocks flag"),
+            arg!(--"py-compat" "Enables improved Python compatibility.  Enabling \
+                this adds methods such as dict.keys and some others."),
             arg!(-s --syntax <PAIR>... "Changes a syntax feature (feature=value) \
                 [possible features: block-start, block-end, variable-start, variable-end, \
                 comment-start, comment-end, line-statement-prefix, \
diff --git a/minijinja-cli/src/main.rs b/minijinja-cli/src/main.rs
index 9ac25a5..fbd3e16 100644
--- a/minijinja-cli/src/main.rs
+++ b/minijinja-cli/src/main.rs
@@ -225,6 +225,9 @@ fn create_env(
     }
 
     minijinja_contrib::add_to_environment(&mut env);
+    if matches.get_flag("py-compat") {
+        env.set_unknown_method_callback(minijinja_contrib::pycompat::unknown_method_callback);
+    }
 
     if matches.get_flag("env") {
         env.add_global("ENV", Value::from_iter(std::env::vars()));
diff --git a/minijinja-contrib/Cargo.toml b/minijinja-contrib/Cargo.toml
index ea432d8..6da7e7d 100644
--- a/minijinja-contrib/Cargo.toml
+++ b/minijinja-contrib/Cargo.toml
@@ -17,6 +17,7 @@ all-features = true
 
 [features]
 default = []
+pycompat = ["minijinja/builtins"]
 datetime = ["time"]
 timezone = ["time-tz"]
 
diff --git a/minijinja-contrib/src/lib.rs b/minijinja-contrib/src/lib.rs
index 08b3195..e7f2df3 100644
--- a/minijinja-contrib/src/lib.rs
+++ b/minijinja-contrib/src/lib.rs
@@ -14,6 +14,10 @@
 
 use minijinja::Environment;
 
+/// Implements Python methods for better compatibility.
+#[cfg(feature = "pycompat")]
+pub mod pycompat;
+
 /// Utility filters.
 pub mod filters;
 
@@ -24,6 +28,9 @@ pub mod globals;
 ///
 /// All the filters that are available will be added, same with global
 /// functions that exist.
+///
+/// **Note:** the `pycompat` support is intentionally not registered
+/// with the environment.
 pub fn add_to_environment(env: &mut Environment) {
     env.add_filter("pluralize", filters::pluralize);
     #[cfg(feature = "datetime")]
diff --git a/minijinja-contrib/src/pycompat.rs b/minijinja-contrib/src/pycompat.rs
new file mode 100644
index 0000000..8d60b9f
--- /dev/null
+++ b/minijinja-contrib/src/pycompat.rs
@@ -0,0 +1,206 @@
+use minijinja::value::{from_args, ValueKind};
+use minijinja::{Error, ErrorKind, State, Value};
+
+/// An unknown method callback implementing python methods on primitives.
+///
+/// This implements a lot of Python methods on basic types so that the
+/// compatibility with Jinja2 templates improves.
+///
+/// ```
+/// use minijinja::Environment;
+/// use minijinja_contrib::pycompat::unknown_method_callback;
+///
+/// let mut env = Environment::new();
+/// env.set_unknown_method_callback(unknown_method_callback);
+/// ```
+///
+/// Today the following methods are implemented:
+///
+/// * `dict.get`
+/// * `dict.items`
+/// * `dict.keys`
+/// * `dict.values`
+/// * `list.count`
+/// * `str.capitalize`
+/// * `str.count`
+/// * `str.find`
+/// * `str.islower`
+/// * `str.isupper`
+/// * `str.lower`
+/// * `str.lstrip`
+/// * `str.replace`
+/// * `str.rstrip`
+/// * `str.strip`
+/// * `str.title`
+/// * `str.upper`
+#[cfg_attr(docsrs, doc(cfg(feature = "pycompat")))]
+pub fn unknown_method_callback(
+    _state: &State,
+    value: &Value,
+    method: &str,
+    args: &[Value],
+) -> Result<Value, Error> {
+    match value.kind() {
+        ValueKind::String => string_methods(value, method, args),
+        ValueKind::Map => map_methods(value, method, args),
+        ValueKind::Seq => seq_methods(value, method, args),
+        _ => Err(Error::from(ErrorKind::UnknownMethod)),
+    }
+}
+
+fn string_methods(value: &Value, method: &str, args: &[Value]) -> Result<Value, Error> {
+    let s = match value.as_str() {
+        Some(s) => s,
+        None => return Err(Error::from(ErrorKind::UnknownMethod)),
+    };
+
+    match method {
+        "upper" => {
+            from_args(args)?;
+            Ok(Value::from(s.to_uppercase()))
+        }
+        "lower" => {
+            from_args(args)?;
+            Ok(Value::from(s.to_lowercase()))
+        }
+        "islower" => {
+            from_args(args)?;
+            Ok(Value::from(s.chars().all(|x| x.is_lowercase())))
+        }
+        "isupper" => {
+            from_args(args)?;
+            Ok(Value::from(s.chars().all(|x| x.is_uppercase())))
+        }
+        "isspace" => {
+            from_args(args)?;
+            Ok(Value::from(s.chars().all(|x| x.is_whitespace())))
+        }
+        "strip" => {
+            let (chars,): (Option<&str>,) = from_args(args)?;
+            Ok(Value::from(if let Some(chars) = chars {
+                s.trim_matches(&chars.chars().collect::<Vec<_>>()[..])
+            } else {
+                s.trim()
+            }))
+        }
+        "lstrip" => {
+            let (chars,): (Option<&str>,) = from_args(args)?;
+            Ok(Value::from(if let Some(chars) = chars {
+                s.trim_start_matches(&chars.chars().collect::<Vec<_>>()[..])
+            } else {
+                s.trim_start()
+            }))
+        }
+        "rstrip" => {
+            let (chars,): (Option<&str>,) = from_args(args)?;
+            Ok(Value::from(if let Some(chars) = chars {
+                s.trim_end_matches(&chars.chars().collect::<Vec<_>>()[..])
+            } else {
+                s.trim_end()
+            }))
+        }
+        "replace" => {
+            let (old, new, count): (&str, &str, Option<i32>) = from_args(args)?;
+            let count = count.unwrap_or(-1);
+            Ok(Value::from(if count < 0 {
+                s.replace(old, new)
+            } else {
+                s.replacen(old, new, count as usize)
+            }))
+        }
+        "title" => {
+            from_args(args)?;
+            // one shall not call into these filters.  However we consider ourselves
+            // privileged.
+            Ok(Value::from(minijinja::filters::title(s.into())))
+        }
+        "capitalize" => {
+            from_args(args)?;
+            // one shall not call into these filters.  However we consider ourselves
+            // privileged.
+            Ok(Value::from(minijinja::filters::capitalize(s.into())))
+        }
+        "count" => {
+            let (what,): (&str,) = from_args(args)?;
+            let mut c = 0;
+            let mut rest = s;
+            while let Some(offset) = rest.find(what) {
+                c += 1;
+                rest = &rest[offset + what.len()..];
+            }
+            Ok(Value::from(c))
+        }
+        "find" => {
+            let (what,): (&str,) = from_args(args)?;
+            Ok(Value::from(match s.find(what) {
+                Some(x) => x as i64,
+                None => -1,
+            }))
+        }
+        _ => Err(Error::from(ErrorKind::UnknownMethod)),
+    }
+}
+
+fn map_methods(value: &Value, method: &str, args: &[Value]) -> Result<Value, Error> {
+    let obj = match value.as_object() {
+        Some(obj) => obj,
+        None => return Err(Error::from(ErrorKind::UnknownMethod)),
+    };
+
+    match method {
+        "keys" => {
+            from_args(args)?;
+            Ok(Value::make_object_iterable(obj.clone(), |obj| {
+                match obj.try_iter() {
+                    Some(iter) => iter,
+                    None => Box::new(None.into_iter()),
+                }
+            }))
+        }
+        "values" => {
+            from_args(args)?;
+            Ok(Value::make_object_iterable(obj.clone(), |obj| {
+                match obj.try_iter_pairs() {
+                    Some(iter) => Box::new(iter.map(|(_, v)| v)),
+                    None => Box::new(None.into_iter()),
+                }
+            }))
+        }
+        "items" => {
+            from_args(args)?;
+            Ok(Value::make_object_iterable(obj.clone(), |obj| {
+                match obj.try_iter_pairs() {
+                    Some(iter) => Box::new(iter.map(|(k, v)| Value::from(vec![k, v]))),
+                    None => Box::new(None.into_iter()),
+                }
+            }))
+        }
+        "get" => {
+            let (key,): (&Value,) = from_args(args)?;
+            Ok(match obj.get_value(key) {
+                Some(value) => value,
+                None => Value::from(()),
+            })
+        }
+        _ => Err(Error::from(ErrorKind::UnknownMethod)),
+    }
+}
+
+fn seq_methods(value: &Value, method: &str, args: &[Value]) -> Result<Value, Error> {
+    let obj = match value.as_object() {
+        Some(obj) => obj,
+        None => return Err(Error::from(ErrorKind::UnknownMethod)),
+    };
+
+    match method {
+        "count" => {
+            let (what,): (&Value,) = from_args(args)?;
+            Ok(Value::from(if let Some(iter) = obj.try_iter() {
+                iter.filter(|x| x == what).count()
+            } else {
+                0
+            }))
+        }
+        _ => Err(Error::from(ErrorKind::UnknownMethod)),
+    }
+}


Test patch:
diff --git a/minijinja-contrib/tests/pycompat.rs b/minijinja-contrib/tests/pycompat.rs
new file mode 100644
index 0000000..dbcbe93
--- /dev/null
+++ b/minijinja-contrib/tests/pycompat.rs
@@ -0,0 +1,58 @@
+#![cfg(feature = "pycompat")]
+use minijinja::{Environment, Value};
+use minijinja_contrib::pycompat::unknown_method_callback;
+use similar_asserts::assert_eq;
+
+fn eval_expr(expr: &str) -> Value {
+    let mut env = Environment::new();
+    env.set_unknown_method_callback(unknown_method_callback);
+    env.compile_expression(expr).unwrap().eval(()).unwrap()
+}
+
+#[test]
+fn test_string_methods() {
+    assert_eq!(eval_expr("'foo'.upper()").as_str(), Some("FOO"));
+    assert_eq!(eval_expr("'FoO'.lower()").as_str(), Some("foo"));
+    assert_eq!(eval_expr("' foo '.strip()").as_str(), Some("foo"));
+    assert_eq!(eval_expr("'!foo?!!!'.strip('!?')").as_str(), Some("foo"));
+    assert_eq!(
+        eval_expr("'!!!foo?!!!'.rstrip('!?')").as_str(),
+        Some("!!!foo")
+    );
+    assert_eq!(
+        eval_expr("'!!!foo?!!!'.lstrip('!?')").as_str(),
+        Some("foo?!!!")
+    );
+    assert!(eval_expr("'foobar'.islower()").is_true());
+    assert!(eval_expr("'FOOBAR'.isupper()").is_true());
+    assert!(eval_expr("' \\n'.isspace()").is_true());
+    assert_eq!(
+        eval_expr("'foobar'.replace('o', 'x')").as_str(),
+        Some("fxxbar")
+    );
+    assert_eq!(
+        eval_expr("'foobar'.replace('o', 'x', 1)").as_str(),
+        Some("fxobar")
+    );
+    assert_eq!(eval_expr("'foo bar'.title()").as_str(), Some("Foo Bar"));
+    assert_eq!(
+        eval_expr("'foo bar'.capitalize()").as_str(),
+        Some("Foo bar")
+    );
+    assert_eq!(eval_expr("'foo barooo'.count('oo')").as_usize(), Some(2));
+    assert_eq!(eval_expr("'foo barooo'.find('oo')").as_usize(), Some(1));
+}
+
+#[test]
+fn test_dict_methods() {
+    assert!(eval_expr("{'x': 42}.keys()|list == ['x']").is_true());
+    assert!(eval_expr("{'x': 42}.values()|list == [42]").is_true());
+    assert!(eval_expr("{'x': 42}.items()|list == [('x', 42)]").is_true());
+    assert!(eval_expr("{'x': 42}.get('x') == 42").is_true());
+    assert!(eval_expr("{'x': 42}.get('y') is none").is_true());
+}
+
+#[test]
+fn test_list_methods() {
+    assert!(eval_expr("[1, 2, 2, 3].count(2) == 2").is_true());
+}