Repository: mitsuhiko/minijinja (rust)

Issue:
`{% import "foo.html" as foo %}` doesn't work as in Jinja
## Description

In Jinja, I can import a partial using the `as foo` syntax, and then print the template using `{{ foo }}`. Minijinja however just prints `{}`. 

## Reproduction steps

Runnable reproduction:

``` rust
use minijinja::Environment;

fn main() {
    let mut env = Environment::new();
    env.add_template("main.html", "{% import \"hello.html\" as hello %} {{ hello }}").unwrap();
    env.add_template("hello.html", "<p>Hello world!</p>").unwrap();

    let template = env.get_template("main.html").unwrap();
    let html = template.render(&{}).unwrap();

    println!("{}", html);
}
```

Additional helpful information:

- Version of minijinja: 2.9.0
- Version of rustc: 1.86
- Operating system and version: macOS

## What did you expect

Program prints `<p>Hello world!</p>`

## Why is this relevant?

We use this as a hack to inline SVG icons into `style` tags in Material for MkDocs:

``` jinja
{% if config.theme.icon.admonition %}
  {% set _ = namespace(style = "\x3cstyle\x3e:root{") %}
  {% for type, icon in config.theme.icon.admonition | items %}
    {% import ".icons/" ~ icon ~ ".svg" as icon %}
    {% set _.style = _.style ~
      "--md-admonition-icon--" ~ type ~ ":" ~
      "url('data:image/svg+xml;charset=utf-8," ~
        icon | replace("\n", "") | urlencode ~
      "');"
    %}
  {% endfor %}
  {% set _.style = _.style ~ "}\x3c/style\x3e" %}
  {{ _.style }}
{% endif %}
```

[Source](https://github.com/squidfunk/mkdocs-material/blob/4c0004e16b1d51511fbd3c8537394069f004ecfe/src/templates/partials/icons.html#L24-L37)

Pull request description:
Import captures body
This fixes MiniJinja not capturing the output like in Jinja2 for imports.

This slightly changes behavior as modules now print out their captured body instead of looking like a map.

Fixes #777

Patch:
diff --git a/.vscode/settings.json b/.vscode/settings.json
index c0b3e8a..49c6fb4 100644
--- a/.vscode/settings.json
+++ b/.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"
 }
\ No newline at end of file
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 69a4cbc..840f746 100644
--- a/CHANGELOG.md
+++ b/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
 
diff --git a/minijinja-cli/examples/include.txt b/minijinja-cli/examples/include.txt
index 2946cda..094067d 100644
--- a/minijinja-cli/examples/include.txt
+++ b/minijinja-cli/examples/include.txt
@@ -1,1 +1,2 @@
-Text from an included template
\ No newline at end of file
+Text from an included template
+{% set x = 42 %}
diff --git a/minijinja/src/compiler/codegen.rs b/minijinja/src/compiler/codegen.rs
index 5b84a9a..f9079b2 100644
--- a/minijinja/src/compiler/codegen.rs
+++ b/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) => {
diff --git a/minijinja/src/vm/mod.rs b/minijinja/src/vm/mod.rs
index 59c8827..70033d0 100644
--- a/minijinja/src/vm/mod.rs
+++ b/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) => {
diff --git a/minijinja/src/vm/module_object.rs b/minijinja/src/vm/module_object.rs
new file mode 100644
index 0000000..c7a616a
--- /dev/null
+++ b/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)
+    }
+}


Test patch:
diff --git a/minijinja/tests/inputs/import_all.txt b/minijinja/tests/inputs/import_all.txt
new file mode 100644
index 0000000..6109396
--- /dev/null
+++ b/minijinja/tests/inputs/import_all.txt
@@ -0,0 +1,7 @@
+{
+}
+---
+{% set x = 23 %}
+{% import 'import-mixed.txt' as module %}
+{{ module|tojson }}
+|||{{ module }}|||
diff --git a/minijinja/tests/inputs/macro_import2.txt b/minijinja/tests/inputs/macro_import2.txt
index 429d7d0..2c3a632 100644
--- a/minijinja/tests/inputs/macro_import2.txt
+++ b/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") }}
\ No newline at end of file
diff --git a/minijinja/tests/inputs/refs/import-mixed.txt b/minijinja/tests/inputs/refs/import-mixed.txt
new file mode 100644
index 0000000..55c3e26
--- /dev/null
+++ b/minijinja/tests/inputs/refs/import-mixed.txt
@@ -0,0 +1,2 @@
+This is the body {{ x }}
+{%- set output = 42 %}
\ No newline at end of file
diff --git a/minijinja/tests/snapshots/test_templates__vm@debug.txt.snap b/minijinja/tests/snapshots/test_templates__vm@debug.txt.snap
index adb47ce..e2049c0 100644
--- a/minijinja/tests/snapshots/test_templates__vm@debug.txt.snap
+++ b/minijinja/tests/snapshots/test_templates__vm@debug.txt.snap
@@ -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",
diff --git a/minijinja/tests/snapshots/test_templates__vm@import_all.txt.snap b/minijinja/tests/snapshots/test_templates__vm@import_all.txt.snap
new file mode 100644
index 0000000..04c8363
--- /dev/null
+++ b/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|||