Import captures body (#778)

Author: mitsuhiko

.vscode/settings.json

@@ -1,6 +1,6 @@
 {
     "rust-analyzer.cargo.allFeatures": true,
-    "rust-analyzer.checkOnSave.command": "clippy",
     "python.formatting.provider": "black",
-    "python.REPL.enableREPLSmartSend": false
+    "python.REPL.enableREPLSmartSend": false,
+    "rust-analyzer.check.command": "clippy"
 }

CHANGELOG.md

@@ -6,6 +6,10 @@ All notable changes to MiniJinja are documented here.
 
 - Fix incorrect permissions when `--output` is used in the CLI.  #772
 - Added `mj_err_get_debug_info` to the C-ABI.  #775
+- Modules now capture their output like they do in Jinja2.  This
+  means that if you do `{% import 'template.j2' as x %}` and you
+  then render `{{ x }}` the output of `template.j2` is rendered as
+  if it was included.  #778
 
 ## 2.9.0
 

minijinja-cli/examples/include.txt

@@ -1 +1,2 @@
-Text from an included template
+Text from an included template
+{% set x = 42 %}

minijinja/src/compiler/codegen.rs

@@ -325,14 +325,14 @@ impl<'source> CodeGenerator<'source> {
             }
             #[cfg(feature = "multi_template")]
             ast::Stmt::Import(import) => {
-                self.add(Instruction::BeginCapture(CaptureMode::Discard));
+                self.add(Instruction::BeginCapture(CaptureMode::Capture));
                 self.add(Instruction::PushWith);
                 self.compile_expr(&import.expr);
                 self.add_with_span(Instruction::Include(false), import.span());
+                self.add(Instruction::EndCapture);
                 self.add(Instruction::ExportLocals);
                 self.add(Instruction::PopFrame);
                 self.compile_assignment(&import.name);
-                self.add(Instruction::EndCapture);
             }
             #[cfg(feature = "multi_template")]
             ast::Stmt::FromImport(from_import) => {

minijinja/src/vm/mod.rs

@@ -31,6 +31,8 @@ mod fuel;
 mod loop_object;
 #[cfg(feature = "macros")]
 mod macro_object;
+#[cfg(feature = "multi_template")]
+mod module_object;
 mod state;
 
 // the cost of a single include against the stack limit.
@@ -664,12 +666,15 @@ impl<'env> Vm<'env> {
                 }
                 #[cfg(feature = "multi_template")]
                 Instruction::ExportLocals => {
+                    let captured = stack.pop();
                     let locals = state.ctx.current_locals_mut();
-                    let mut module = value_map_with_capacity(locals.len());
+                    let mut values = value_map_with_capacity(locals.len());
                     for (key, value) in locals.iter() {
-                        module.insert(Value::from(*key), value.clone());
+                        values.insert(Value::from(*key), value.clone());
                     }
-                    stack.push(Value::from_object(module));
+                    stack.push(Value::from_object(crate::vm::module_object::Module::new(
+                        values, captured,
+                    )));
                 }
                 #[cfg(feature = "multi_template")]
                 Instruction::CallBlock(name) => {

minijinja/src/vm/module_object.rs

@@ -0,0 +1,32 @@
+use std::fmt;
+use std::sync::Arc;
+
+use crate::value::{Enumerator, Object, Value, ValueMap};
+
+/// Holds the locals of a module.
+#[derive(Debug, Default)]
+pub(crate) struct Module {
+    values: ValueMap,
+    captured: Value,
+}
+
+impl Module {
+    pub fn new(values: ValueMap, captured: Value) -> Self {
+        Self { values, captured }
+    }
+}
+
+impl Object for Module {
+    fn get_value(self: &Arc<Self>, key: &Value) -> Option<Value> {
+        self.values.get(key).cloned()
+    }
+
+    fn enumerate(self: &Arc<Self>) -> Enumerator {
+        let keys = self.values.keys().cloned();
+        Enumerator::Values(keys.collect())
+    }
+
+    fn render(self: &Arc<Self>, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "{}", self.captured)
+    }
+}

minijinja/tests/inputs/import_all.txt

@@ -0,0 +1,7 @@
+{
+}
+---
+{% set x = 23 %}
+{% import 'import-mixed.txt' as module %}
+{{ module|tojson }}
+|||{{ module }}|||

minijinja/tests/inputs/macro_import2.txt

@@ -1,7 +1,7 @@
 {"c": "The C Variable"}
 ---
 {% import "include_with_var_and_macro.txt" as helpers -%}
-{{ helpers }}
+{{ dict(helpers) }}
 missing: {{ helpers.missing }}
 title: {{ helpers.title }}
 helper: {{ helpers.helper("a", "b") }}

minijinja/tests/inputs/refs/import-mixed.txt

@@ -0,0 +1,2 @@
+This is the body {{ x }}
+{%- set output = 42 %}

minijinja/tests/snapshots/[email protected]

@@ -122,6 +122,7 @@ State {
             "call_macro.txt",
             "debug.txt",
             "example_macro.txt",
+            "import-mixed.txt",
             "include_with_var_and_macro.txt",
             "layout_with_var.txt",
             "self-extends.txt",

minijinja/tests/snapshots/test_templates__vm@import_all.txt.snap

@@ -0,0 +1,8 @@
+---
+source: minijinja/tests/test_templates.rs
+description: "{% set x = 23 %}\n{% import 'import-mixed.txt' as module %}\n{{ module|tojson }}\n|||{{ module }}|||"
+info: {}
+input_file: minijinja/tests/inputs/import_all.txt
+---
+{"output":42}
+|||This is the body 23|||