Support custom simple block tags

src/render/filters.rs

@@ -104,8 +104,8 @@ impl ResolveFilter for AddFilter {
         Ok(match (variable.to_bigint(), right.to_bigint()) {
             (Some(variable), Some(right)) => Some(Content::Int(variable + right)),
             _ => {
-                let variable = variable.to_py(py)?;
-                let right = right.to_py(py)?;
+                let variable = variable.to_py(py);
+                let right = right.to_py(py);
                 variable.add(right).ok().map(Content::Py)
             }
         })

src/render/tags.rs

@@ -12,7 +12,7 @@ use pyo3::types::{PyBool, PyDict, PyList, PyNone, PyString, PyTuple};
 use super::types::{AsBorrowedContent, Content, Context, PyContext};
 use super::{Evaluate, Render, RenderResult, Resolve, ResolveFailures, ResolveResult};
 use crate::error::{AnnotatePyErr, PyRenderError, RenderError};
-use crate::parse::{For, IfCondition, SimpleTag, Tag, Url};
+use crate::parse::{For, IfCondition, SimpleBlockTag, SimpleTag, Tag, TagElement, Url};
 use crate::template::django_rusty_templates::NoReverseMatch;
 use crate::types::TemplateString;
 use crate::utils::PyResultMethods;
@@ -420,7 +420,7 @@ impl Contains<Option<Content<'_, '_>>> for Content<'_, '_> {
                 _ => None,
             },
             Some(Content::Py(other)) => {
-                let obj = self.to_py(other.py()).ok()?;
+                let obj = self.to_py(other.py());
                 obj.contains(other).ok()
             }
             Some(Content::String(other)) => match self {
@@ -646,6 +646,7 @@ impl Render for Tag {
             Self::For(for_tag) => for_tag.render(py, template, context)?,
             Self::Load => Cow::Borrowed(""),
             Self::SimpleTag(simple_tag) => simple_tag.render(py, template, context)?,
+            Self::SimpleBlockTag(simple_tag) => simple_tag.render(py, template, context)?,
             Self::Url(url) => url.render(py, template, context)?,
         })
     }
@@ -756,86 +757,175 @@ impl Render for For {
     }
 }
 
-impl SimpleTag {
-    fn call_tag<'t>(
+fn call_tag<'t>(
+    py: Python<'_>,
+    func: &Arc<Py<PyAny>>,
+    at: (usize, usize),
+    template: TemplateString<'t>,
+    args: VecDeque<Bound<'_, PyAny>>,
+    kwargs: Bound<'_, PyDict>,
+) -> RenderResult<'t> {
+    let func = func.bind(py);
+    match func.call(
+        PyTuple::new(py, args).expect("All arguments should be valid Python objects"),
+        Some(&kwargs),
+    ) {
+        Ok(content) => Ok(Cow::Owned(content.to_string())),
+        Err(error) => Err(error.annotate(py, at, "here", template).into()),
+    }
+}
+
+fn add_context_to_args<'py>(
+    py: Python<'py>,
+    args: &mut VecDeque<Bound<'py, PyAny>>,
+    context: &mut Context,
+) -> Result<Bound<'py, PyAny>, PyErr> {
+    // Take ownership of `context` so we can pass it to Python.
+    // The `context` variable now points to an empty `Context` instance which will not be
+    // used except as a placeholder.
+    let swapped_context = std::mem::take(context);
+
+    // Wrap the context as a Python object and add it to the call args
+    let py_context = Bound::new(py, PyContext::new(swapped_context))?.into_any();
+    args.push_front(py_context.clone());
+
+    Ok(py_context)
+}
+
+fn retrieve_context<'py>(py: Python<'py>, py_context: Bound<'py, PyAny>, context: &mut Context) {
+    // Retrieve the PyContext wrapper from Python
+    let extracted_context: PyContext = py_context
+        .extract()
+        .expect("The type of py_context should not have changed");
+    // Ensure we only hold one reference in Rust by dropping the Python object.
+    drop(py_context);
+
+    // Try to remove the Context from the PyContext
+    let inner_context = match Arc::try_unwrap(extracted_context.context) {
+        // Fast path when we have the only reference in the Arc.
+        Ok(inner_context) => inner_context
+            .into_inner()
+            .expect("Mutex should be unlocked because Arc refcount is one."),
+        // Slow path when Python has held on to the context for some reason.
+        // We can still do the right thing by cloning.
+        Err(inner_context) => {
+            let guard = inner_context
+                .lock_py_attached(py)
+                .expect("Mutex should not be poisoned");
+            guard.clone_ref(py)
+        }
+    };
+    // Put the Context back in `context`
+    let _ = std::mem::replace(context, inner_context);
+}
+
+fn build_arg<'py>(
+    py: Python<'py>,
+    template: TemplateString<'_>,
+    context: &mut Context,
+    arg: &TagElement,
+) -> Result<Bound<'py, PyAny>, PyRenderError> {
+    let arg = match arg.resolve(py, template, context, ResolveFailures::Raise)? {
+        Some(arg) => arg.to_py(py),
+        None => PyString::intern(py, "").into_any(),
+    };
+    Ok(arg)
+}
+
+fn build_args<'py>(
+    py: Python<'py>,
+    template: TemplateString<'_>,
+    context: &mut Context,
+    args: &[TagElement],
+) -> Result<VecDeque<Bound<'py, PyAny>>, PyRenderError> {
+    args.iter()
+        .map(|arg| build_arg(py, template, context, arg))
+        .collect()
+}
+
+fn build_kwargs<'py>(
+    py: Python<'py>,
+    template: TemplateString<'_>,
+    context: &mut Context,
+    kwargs: &Vec<(String, TagElement)>,
+) -> Result<Bound<'py, PyDict>, PyRenderError> {
+    let _kwargs = PyDict::new(py);
+    for (key, value) in kwargs {
+        let value = value.resolve(py, template, context, ResolveFailures::Raise)?;
+        _kwargs.set_item(key, value)?;
+    }
+    Ok(_kwargs)
+}
+
+fn store_target_var<'t>(
+    py: Python<'_>,
+    context: &mut Context,
+    content: Cow<'t, str>,
+    target_var: &Option<String>,
+) -> Cow<'t, str> {
+    match target_var {
+        None => content,
+        Some(target_var) => {
+            let content = PyString::new(py, &content).into_any();
+            context.insert(target_var.clone(), content);
+            Cow::Borrowed("")
+        }
+    }
+}
+
+impl Render for SimpleTag {
+    fn render<'t>(
         &self,
         py: Python<'_>,
         template: TemplateString<'t>,
-        args: VecDeque<Bound<'_, PyAny>>,
-        kwargs: Bound<'_, PyDict>,
+        context: &mut Context,
     ) -> RenderResult<'t> {
-        let func = self.func.bind(py);
-        match func.call(
-            PyTuple::new(py, args).expect("All arguments should be valid Python objects"),
-            Some(&kwargs),
-        ) {
-            Ok(content) => Ok(Cow::Owned(content.to_string())),
-            Err(error) => Err(error.annotate(py, self.at, "here", template).into()),
-        }
+        let mut args = build_args(py, template, context, &self.args)?;
+        let kwargs = build_kwargs(py, template, context, &self.kwargs)?;
+        let content = if self.takes_context {
+            let py_context = add_context_to_args(py, &mut args, context)?;
+
+            // Actually call the tag
+            let result = call_tag(py, &self.func, self.at, template, args, kwargs);
+
+            retrieve_context(py, py_context, context);
+
+            // Return the result of calling the tag
+            result?
+        } else {
+            call_tag(py, &self.func, self.at, template, args, kwargs)?
+        };
+        Ok(store_target_var(py, context, content, &self.target_var))
     }
 }
 
-impl Render for SimpleTag {
+impl Render for SimpleBlockTag {
     fn render<'t>(
         &self,
         py: Python<'_>,
         template: TemplateString<'t>,
         context: &mut Context,
     ) -> RenderResult<'t> {
-        let mut args = VecDeque::new();
-        for arg in &self.args {
-            match arg.resolve(py, template, context, ResolveFailures::Raise)? {
-                None => return Ok(Cow::Borrowed("")),
-                Some(arg) => args.push_back(arg.to_py(py)?),
-            }
-        }
-        let kwargs = PyDict::new(py);
-        for (key, value) in &self.kwargs {
-            let value = value.resolve(py, template, context, ResolveFailures::Raise)?;
-            kwargs.set_item(key, value)?;
-        }
-        if self.takes_context {
-            // Take ownership of `context` so we can pass it to Python.
-            // The `context` variable now points to an empty `Context` instance which will not be
-            // used except as a placeholder.
-            let swapped_context = std::mem::take(context);
+        let mut args = build_args(py, template, context, &self.args)?;
+        let kwargs = build_kwargs(py, template, context, &self.kwargs)?;
+
+        let content = self.nodes.render(py, template, context)?;
+        let content = PyString::new(py, &content).into_any();
+        args.push_front(content);
 
-            // Wrap the context as a Python object and add it to the call args
-            let py_context = Bound::new(py, PyContext::new(swapped_context))?.into_any();
-            args.push_front(py_context.clone());
+        let content = if self.takes_context {
+            let py_context = add_context_to_args(py, &mut args, context)?;
 
             // Actually call the tag
-            let result = self.call_tag(py, template, args, kwargs);
-
-            // Retrieve the PyContext wrapper from Python
-            let extracted_context: PyContext = py_context
-                .extract()
-                .expect("The type of py_context should not have changed");
-            // Ensure we only hold one reference in Rust by dropping the Python object.
-            drop(py_context);
-
-            // Try to remove the Context from the PyContext
-            let inner_context = match Arc::try_unwrap(extracted_context.context) {
-                // Fast path when we have the only reference in the Arc.
-                Ok(inner_context) => inner_context
-                    .into_inner()
-                    .expect("Mutex should be unlocked because Arc refcount is one."),
-                // Slow path when Python has held on to the context for some reason.
-                // We can still do the right thing by cloning.
-                Err(inner_context) => {
-                    let guard = inner_context
-                        .lock_py_attached(py)
-                        .expect("Mutex should not be poisoned");
-                    guard.clone_ref(py)
-                }
-            };
-            // Put the Context back in `context`
-            let _ = std::mem::replace(context, inner_context);
+            let result = call_tag(py, &self.func, self.at, template, args, kwargs);
+
+            retrieve_context(py, py_context, context);
 
             // Return the result of calling the tag
-            result
+            result?
         } else {
-            self.call_tag(py, template, args, kwargs)
-        }
+            call_tag(py, &self.func, self.at, template, args, kwargs)?
+        };
+        Ok(store_target_var(py, context, content, &self.target_var))
     }
 }

src/render/types.rs

@@ -448,8 +448,8 @@ impl<'t, 'py> Content<'t, 'py> {
         }
     }
 
-    pub fn to_py(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
-        Ok(match self {
+    pub fn to_py(&self, py: Python<'py>) -> Bound<'py, PyAny> {
+        match self {
             Self::Py(object) => object.clone(),
             Self::Int(i) => i
                 .into_pyobject(py)
@@ -472,13 +472,19 @@ impl<'t, 'py> Content<'t, 'py> {
                     let string = s
                         .into_pyobject(py)
                         .expect("A string can always be converted to a Python str.");
-                    let safestring = py.import(intern!(py, "django.utils.safestring"))?;
-                    let mark_safe = safestring.getattr(intern!(py, "mark_safe"))?;
-                    mark_safe.call1((string,))?
+                    let safestring = py
+                        .import(intern!(py, "django.utils.safestring"))
+                        .expect("Should be able to import `django.utils.safestring`");
+                    let mark_safe = safestring
+                        .getattr(intern!(py, "mark_safe"))
+                        .expect("`safestring` should have a `mark_safe` function");
+                    mark_safe
+                        .call1((string,))
+                        .expect("`mark_safe` should not raise if given a string")
                 }
             },
             Self::Bool(b) => PyBool::new(py, *b).to_owned().into_any(),
-        })
+        }
     }
 }
 

tests/tags/test_custom.py → tests/tags/test_custom_simple.py

@@ -38,6 +38,35 @@ def test_simple_tag_double_missing_variable():
     assert rust_template.render({}) == ""
 
 
+def test_simple_tag_multiply_missing_variables():
+    template = "{% load multiply from custom_tags %}{% multiply foo bar eggs %}"
+
+    django_template = engines["django"].from_string(template)
+    rust_template = engines["rusty"].from_string(template)
+
+    with pytest.raises(TypeError) as exc_info:
+        django_template.render({})
+
+    error = str(exc_info.value)
+    assert error == "can't multiply sequence by non-int of type 'str'"
+
+    with pytest.raises(TypeError) as exc_info:
+        rust_template.render({})
+
+    error = str(exc_info.value)
+    assert (
+        error
+        == """\
+  × can't multiply sequence by non-int of type 'str'
+   ╭────
+ 1 │ {% load multiply from custom_tags %}{% multiply foo bar eggs %}
+   ·                                     ─────────────┬─────────────
+   ·                                                  ╰── here
+   ╰────
+"""
+    )
+
+
 def test_simple_tag_kwargs():
     template = "{% load table from custom_tags %}{% table foo='bar' spam=1 %}"
 
@@ -59,13 +88,13 @@ def test_simple_tag_positional_and_kwargs():
 
 
 def test_simple_tag_double_as_variable():
-    template = "{% load double from custom_tags %}{% double 3 as foo %}{{ foo }}"
+    template = "{% load double from custom_tags %}{% double 3 as foo %}{{ foo }}{{ foo }}"
 
     django_template = engines["django"].from_string(template)
     rust_template = engines["rusty"].from_string(template)
 
-    assert django_template.render({}) == "6"
-    assert rust_template.render({}) == "6"
+    assert django_template.render({}) == "66"
+    assert rust_template.render({}) == "66"
 
 
 def test_simple_tag_double_kwarg_as_variable():