Merge pull request #255 from Kijewski/pr-enum

Implement `enum` variants

.github/workflows/rust.yml

@@ -172,7 +172,8 @@ jobs:
         with:
           tool: cargo-nextest
       - uses: Swatinem/rust-cache@v2
-      - run: cd ${{ matrix.package }} && cargo nextest run --no-tests=warn
+      - run: cd ${{ matrix.package }} && cargo build --all-targets
+      - run: cd ${{ matrix.package }} && cargo nextest run --all-targets --no-fail-fast --no-tests=warn
       - run: cd ${{ matrix.package }} && cargo clippy --all-targets -- -D warnings
 
   MSRV:

book/src/creating_templates.md

@@ -118,6 +118,85 @@ recognized:
   struct HelloTemplate<'a> { ... }
   ```
 
+## Templating `enum`s
+
+You can add derive `Template`s for `struct`s and `enum`s.
+If you add `#[template()]` only to the item itself, both item kinds work exactly the same.
+But with `enum`s you also have the option to add a specialized implementation to one, some,
+or all variants:
+
+```rust
+#[derive(Debug, Template)]
+#[template(path = "area.txt")]
+enum Area {
+    Square(f32),
+    Rectangle { a: f32, b: f32 },
+    Circle { radius: f32 },
+}
+```
+
+```jinja2
+{%- match self -%}
+    {%- when Self::Square(side) -%}
+        {{side}}^2
+    {%- when Self::Rectangle { a, b} -%}
+        {{a}} * {{b}}
+    {%- when Self::Circle { radius } -%}
+        pi * {{radius}}^2
+{%- endmatch -%}
+```
+
+will give you the same results as:
+
+```rust
+#[derive(Template, Debug)]
+#[template(ext = "txt")]
+enum AreaPerVariant {
+    #[template(source = "{{self.0}}^2")]
+    Square(f32),
+    #[template(source = "{{a}} * {{b}}")]
+    Rectangle { a: f32, b: f32 },
+    #[template(source = "pi * {{radius}}^2")]
+    Circle { radius: f32 },
+}
+```
+
+As you can see with the `ext` attribute, `enum` variants inherit most settings of the `enum`:
+`config`, `escape`, `ext`, `syntax`, and `whitespace`.
+Not inherited are: `block`, and `print`.
+
+If there is no `#[template]` annotation for an `enum` variant,
+then the `enum` needs a default implementation, which will be used if `self` is this variant.
+A good compromise between annotating only the template, or all its variants,
+might be using the `block` argument on the members:
+
+```rust
+#[derive(Template, Debug)]
+#[template(path = "area.txt")]
+enum AreaWithBlocks {
+    #[template(block = "square")]
+    Square(f32),
+    #[template(block = "rectangle")]
+    Rectangle { a: f32, b: f32 },
+    #[template(block = "circle")]
+    Circle { radius: f32 },
+}
+```
+
+```jinja2
+{%- block square -%}
+    {{self.0}}^2
+{%- endblock -%}
+
+{%- block rectangle -%}
+    {{a}} * {{b}}
+{%- endblock -%}
+
+{%- block circle -%}
+    pi * {{radius}}^2
+{%- endblock -%}
+```
+
 ## Documentation as template code
 [#documentation-as-template-code]: #documentation-as-template-code
 

rinja/src/helpers.rs

@@ -269,3 +269,11 @@ impl<L: FastWritable, R: FastWritable> FastWritable for Concat<L, R> {
         self.1.write_into(dest)
     }
 }
+
+pub trait EnumVariantTemplate {
+    fn render_into_with_values<W: fmt::Write + ?Sized>(
+        &self,
+        writer: &mut W,
+        values: &dyn crate::Values,
+    ) -> crate::Result<()>;
+}

rinja_derive/src/generator.rs

@@ -25,9 +25,12 @@ pub(crate) fn template_to_string(
     input: &TemplateInput<'_>,
     contexts: &HashMap<&Arc<Path>, Context<'_>, FxBuildHasher>,
     heritage: Option<&Heritage<'_, '_>>,
-    target: Option<&str>,
+    tmpl_kind: TmplKind,
 ) -> Result<usize, CompileError> {
-    let ctx = &contexts[&input.path];
+    if tmpl_kind == TmplKind::Struct {
+        buf.write("const _: () = { extern crate rinja as rinja;");
+    }
+
     let generator = Generator::new(
         input,
         contexts,
@@ -36,13 +39,27 @@ pub(crate) fn template_to_string(
         input.block.is_some(),
         0,
     );
-    let mut result = generator.build(ctx, buf, target);
-    if let Err(err) = &mut result {
-        if err.span.is_none() {
+    let size_hint = match generator.impl_template(buf, tmpl_kind) {
+        Err(mut err) if err.span.is_none() => {
             err.span = input.source_span;
+            Err(err)
         }
+        result => result,
+    }?;
+
+    if tmpl_kind == TmplKind::Struct {
+        impl_everything(input.ast, buf);
+        buf.write("};");
     }
-    result
+    Ok(size_hint)
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub(crate) enum TmplKind {
+    /// [`rinja::Template`]
+    Struct,
+    /// [`rinja::helpers::EnumVariantTemplate`]
+    Variant,
 }
 
 struct Generator<'a, 'h> {
@@ -97,31 +114,18 @@ impl<'a, 'h> Generator<'a, 'h> {
         }
     }
 
-    // Takes a Context and generates the relevant implementations.
-    fn build(
-        mut self,
-        ctx: &Context<'a>,
-        buf: &mut Buffer,
-        target: Option<&str>,
-    ) -> Result<usize, CompileError> {
-        if target.is_none() {
-            buf.write("const _: () = { extern crate rinja as rinja;");
-        }
-        let size_hint = self.impl_template(ctx, buf, target.unwrap_or("rinja::Template"))?;
-        if target.is_none() {
-            impl_everything(self.input.ast, buf);
-            buf.write("};");
-        }
-        Ok(size_hint)
-    }
-
     // Implement `Template` for the given context struct.
     fn impl_template(
-        &mut self,
-        ctx: &Context<'a>,
+        mut self,
         buf: &mut Buffer,
-        target: &str,
+        tmpl_kind: TmplKind,
     ) -> Result<usize, CompileError> {
+        let ctx = &self.contexts[&self.input.path];
+
+        let target = match tmpl_kind {
+            TmplKind::Struct => "rinja::Template",
+            TmplKind::Variant => "rinja::helpers::EnumVariantTemplate",
+        };
         write_header(self.input.ast, buf, target);
         buf.write(
             "fn render_into_with_values<RinjaW>(\
@@ -161,12 +165,12 @@ impl<'a, 'h> Generator<'a, 'h> {
 
         let size_hint = self.impl_template_inner(ctx, buf)?;
 
-        buf.write(format_args!(
-            "\
-                rinja::Result::Ok(())\
-            }}\
-            const SIZE_HINT: rinja::helpers::core::primitive::usize = {size_hint}usize;",
-        ));
+        buf.write("rinja::Result::Ok(()) }");
+        if tmpl_kind == TmplKind::Struct {
+            buf.write(format_args!(
+                "const SIZE_HINT: rinja::helpers::core::primitive::usize = {size_hint}usize;",
+            ));
+        }
 
         buf.write('}');
         Ok(size_hint)

rinja_derive/src/input.rs

@@ -271,6 +271,63 @@ impl TemplateInput<'_> {
     }
 }
 
+pub(crate) enum AnyTemplateArgs {
+    Struct(TemplateArgs),
+    Enum {
+        enum_args: Option<PartialTemplateArgs>,
+        vars_args: Vec<Option<PartialTemplateArgs>>,
+        has_default_impl: bool,
+    },
+}
+
+impl AnyTemplateArgs {
+    pub(crate) fn new(ast: &syn::DeriveInput) -> Result<Self, CompileError> {
+        let syn::Data::Enum(enum_data) = &ast.data else {
+            return Ok(Self::Struct(TemplateArgs::new(ast)?));
+        };
+
+        let enum_args = PartialTemplateArgs::new(ast, &ast.attrs)?;
+        let vars_args = enum_data
+            .variants
+            .iter()
+            .map(|variant| PartialTemplateArgs::new(ast, &variant.attrs))
+            .collect::<Result<Vec<_>, _>>()?;
+        if vars_args.is_empty() {
+            return Ok(Self::Struct(TemplateArgs::from_partial(ast, enum_args)?));
+        }
+
+        let mut needs_default_impl = vars_args.len();
+        let enum_source = enum_args.as_ref().and_then(|v| v.source.as_ref());
+        for (variant, var_args) in enum_data.variants.iter().zip(&vars_args) {
+            if var_args
+                .as_ref()
+                .and_then(|v| v.source.as_ref())
+                .or(enum_source)
+                .is_none()
+            {
+                return Err(CompileError::new_with_span(
+                    #[cfg(not(feature = "code-in-doc"))]
+                    "either all `enum` variants need a `path` or `source` argument, \
+                    or the `enum` itself needs a default implementation",
+                    #[cfg(feature = "code-in-doc")]
+                    "either all `enum` variants need a `path`, `source` or `in_doc` argument, \
+                    or the `enum` itself needs a default implementation",
+                    None,
+                    Some(variant.ident.span()),
+                ));
+            } else if !var_args.is_none() {
+                needs_default_impl -= 1;
+            }
+        }
+
+        Ok(Self::Enum {
+            enum_args,
+            vars_args,
+            has_default_impl: needs_default_impl > 0,
+        })
+    }
+}
+
 #[derive(Debug)]
 pub(crate) struct TemplateArgs {
     pub(crate) source: (Source, Option<Span>),
@@ -626,6 +683,17 @@ pub(crate) enum PartialTemplateArgsSource {
     InDoc(Span, Source),
 }
 
+impl PartialTemplateArgsSource {
+    pub(crate) fn span(&self) -> Span {
+        match self {
+            Self::Path(s) => s.span(),
+            Self::Source(s) => s.span(),
+            #[cfg(feature = "code-in-doc")]
+            Self::InDoc(s, _) => s.span(),
+        }
+    }
+}
+
 // implement PartialTemplateArgs::new()
 const _: () = {
     impl PartialTemplateArgs {