feat: add dynamic file structures in loop using yield-tag (#1855)

Add jinja2 extension for yield tag, allow _render_path to generate multiple paths and contexts when yield tag is used.

The tag is only allowed in path render contexts. When rendering within a file, it doesn't make sense. Use the normal for tag there. If you use yield, you'll get an exception.

Fixes #1271

copier/errors.py

@@ -112,6 +112,14 @@ def __init__(self, features: Sequence[str]):
         )
 
 
+class YieldTagInFileError(CopierError):
+    """A yield tag is used in the file content, but it is not allowed."""
+
+
+class MultipleYieldTagsError(CopierError):
+    """Multiple yield tags are used in one path name, but it is not allowed."""
+
+
 # Warnings
 class CopierWarning(Warning):
     """Base class for all other Copier warnings."""

copier/jinja_ext.py

@@ -0,0 +1,123 @@
+"""Jinja2 extensions built for Copier."""
+
+from __future__ import annotations
+
+from typing import Any, Callable, Iterable
+
+from jinja2 import nodes
+from jinja2.exceptions import UndefinedError
+from jinja2.ext import Extension
+from jinja2.parser import Parser
+from jinja2.sandbox import SandboxedEnvironment
+
+from copier.errors import MultipleYieldTagsError
+
+
+class YieldEnvironment(SandboxedEnvironment):
+    """Jinja2 environment with attributes from the YieldExtension.
+
+    This is simple environment class that extends the SandboxedEnvironment
+    for use with the YieldExtension, mainly for avoiding type errors.
+
+    We use the SandboxedEnvironment because we want to minimize the risk of hidden malware
+    in the templates. Of course we still have the post-copy tasks to worry about, but at least
+    they are more visible to the final user.
+    """
+
+    yield_name: str | None
+    yield_iterable: Iterable[Any] | None
+
+    def __init__(self, *args: Any, **kwargs: Any):
+        super().__init__(*args, **kwargs)
+        self.extend(yield_name=None, yield_iterable=None)
+
+
+class YieldExtension(Extension):
+    """Jinja2 extension for the `yield` tag.
+
+    If `yield` tag is used in a template, this extension sets following attribute to the
+    jinja environment:
+
+    - `yield_name`: The name of the variable that will be yielded.
+    - `yield_iterable`: The variable that will be looped over.
+
+    Note that this extension just sets the attributes but renders templates as usual.
+    It is the caller's responsibility to use the `yield_context` attribute in the template to
+    generate the desired output.
+
+    !!! example
+
+        ```pycon
+        >>> from copier.jinja_ext import YieldEnvironment, YieldExtension
+        >>> env = YieldEnvironment(extensions=[YieldExtension])
+        >>> template = env.from_string("{% yield single_var from looped_var %}{{ single_var }}{% endyield %}")
+        >>> template.render({"looped_var": [1, 2, 3]})
+        ''
+        >>> env.yield_name
+        'single_var'
+        >>> env.yield_iterable
+        [1, 2, 3]
+        ```
+    """
+
+    tags = {"yield"}
+
+    environment: YieldEnvironment
+
+    def preprocess(
+        self, source: str, _name: str | None, _filename: str | None = None
+    ) -> str:
+        """Preprocess hook to reset attributes before rendering."""
+        self.environment.yield_name = self.environment.yield_iterable = None
+
+        return source
+
+    def parse(self, parser: Parser) -> nodes.Node:
+        """Parse the `yield` tag."""
+        lineno = next(parser.stream).lineno
+
+        yield_name: nodes.Name = parser.parse_assign_target(name_only=True)
+        parser.stream.expect("name:from")
+        yield_iterable = parser.parse_expression()
+        body = parser.parse_statements(("name:endyield",), drop_needle=True)
+
+        return nodes.CallBlock(
+            self.call_method(
+                "_yield_support",
+                [nodes.Const(yield_name.name), yield_iterable],
+            ),
+            [],
+            [],
+            body,
+            lineno=lineno,
+        )
+
+    def _yield_support(
+        self, yield_name: str, yield_iterable: Iterable[Any], caller: Callable[[], str]
+    ) -> str:
+        """Support function for the yield tag.
+
+        Sets the `yield_name` and `yield_iterable` attributes in the environment then calls
+        the provided caller function. If an UndefinedError is raised, it returns an empty string.
+        """
+        if (
+            self.environment.yield_name is not None
+            or self.environment.yield_iterable is not None
+        ):
+            raise MultipleYieldTagsError(
+                "Attempted to parse the yield tag twice. Only one yield tag is allowed per path name.\n"
+                f'A yield tag with the name: "{self.environment.yield_name}" and iterable: "{self.environment.yield_iterable}" already exists.'
+            )
+
+        self.environment.yield_name = yield_name
+        self.environment.yield_iterable = yield_iterable
+
+        try:
+            res = caller()
+
+        # expression like `dict.attr` will always raise UndefinedError
+        # so we catch it here and return an empty string
+        except UndefinedError:
+            res = ""
+
+        return res

copier/main.py

@@ -29,7 +29,6 @@
 from unicodedata import normalize
 
 from jinja2.loaders import FileSystemLoader
-from jinja2.sandbox import SandboxedEnvironment
 from pathspec import PathSpec
 from plumbum import ProcessExecutionError, colors
 from plumbum.cli.terminal import ask
@@ -44,7 +43,9 @@
     ExtensionNotFoundError,
     UnsafeTemplateError,
     UserMessageError,
+    YieldTagInFileError,
 )
+from .jinja_ext import YieldEnvironment, YieldExtension
 from .subproject import Subproject
 from .template import Task, Template
 from .tools import (
@@ -541,7 +542,7 @@ def all_exclusions(self) -> Sequence[str]:
         return self.template.exclude + tuple(self.exclude)
 
     @cached_property
-    def jinja_env(self) -> SandboxedEnvironment:
+    def jinja_env(self) -> YieldEnvironment:
         """Return a pre-configured Jinja environment.
 
         Respects template settings.
@@ -550,14 +551,11 @@ def jinja_env(self) -> SandboxedEnvironment:
         loader = FileSystemLoader(paths)
         default_extensions = [
             "jinja2_ansible_filters.AnsibleCoreFiltersExtension",
+            YieldExtension,
         ]
         extensions = default_extensions + list(self.template.jinja_extensions)
-        # We want to minimize the risk of hidden malware in the templates
-        # so we use the SandboxedEnvironment instead of the regular one.
-        # Of course we still have the post-copy tasks to worry about, but at least
-        # they are more visible to the final user.
         try:
-            env = SandboxedEnvironment(
+            env = YieldEnvironment(
                 loader=loader, extensions=extensions, **self.template.envops
             )
         except ModuleNotFoundError as error:
@@ -607,19 +605,25 @@ def _render_template(self) -> None:
         for src in scantree(str(self.template_copy_root), follow_symlinks):
             src_abspath = Path(src.path)
             src_relpath = Path(src_abspath).relative_to(self.template.local_abspath)
-            dst_relpath = self._render_path(
+            dst_relpaths_ctxs = self._render_path(
                 Path(src_abspath).relative_to(self.template_copy_root)
             )
-            if dst_relpath is None or self.match_exclude(dst_relpath):
-                continue
-            if src.is_symlink() and self.template.preserve_symlinks:
-                self._render_symlink(src_relpath, dst_relpath)
-            elif src.is_dir(follow_symlinks=follow_symlinks):
-                self._render_folder(dst_relpath)
-            else:
-                self._render_file(src_relpath, dst_relpath)
+            for dst_relpath, ctx in dst_relpaths_ctxs:
+                if self.match_exclude(dst_relpath):
+                    continue
+                if src.is_symlink() and self.template.preserve_symlinks:
+                    self._render_symlink(src_relpath, dst_relpath)
+                elif src.is_dir(follow_symlinks=follow_symlinks):
+                    self._render_folder(dst_relpath)
+                else:
+                    self._render_file(src_relpath, dst_relpath, extra_context=ctx or {})
 
-    def _render_file(self, src_relpath: Path, dst_relpath: Path) -> None:
+    def _render_file(
+        self,
+        src_relpath: Path,
+        dst_relpath: Path,
+        extra_context: AnyByStrDict | None = None,
+    ) -> None:
         """Render one file.
 
         Args:
@@ -629,6 +633,8 @@ def _render_file(self, src_relpath: Path, dst_relpath: Path) -> None:
             dst_relpath:
                 File to be created. It must be a path relative to the subproject
                 root.
+            extra_context:
+                Additional variables to use for rendering the template.
         """
         # TODO Get from main.render_file()
         assert not src_relpath.is_absolute()
@@ -644,7 +650,13 @@ def _render_file(self, src_relpath: Path, dst_relpath: Path) -> None:
                 # suffix is empty, fallback to copy
                 new_content = src_abspath.read_bytes()
             else:
-                new_content = tpl.render(**self._render_context()).encode()
+                new_content = tpl.render(
+                    **self._render_context(), **(extra_context or {})
+                ).encode()
+                if self.jinja_env.yield_name:
+                    raise YieldTagInFileError(
+                        f"File {src_relpath} contains a yield tag, but it is not allowed."
+                    )
         else:
             new_content = src_abspath.read_bytes()
         dst_abspath = self.subproject.local_abspath / dst_relpath
@@ -716,8 +728,85 @@ def _render_folder(self, dst_relpath: Path) -> None:
             dst_abspath = self.subproject.local_abspath / dst_relpath
             dst_abspath.mkdir(parents=True, exist_ok=True)
 
-    def _render_path(self, relpath: Path) -> Path | None:
-        """Render one relative path.
+    def _adjust_rendered_part(self, rendered_part: str) -> str:
+        """Adjust the rendered part if necessary.
+
+        If `{{ _copier_conf.answers_file }}` becomes the full path,
+        restore part to be just the end leaf.
+
+        Args:
+            rendered_part:
+                The rendered part of the path to adjust.
+
+        """
+        if str(self.answers_relpath) == rendered_part:
+            return Path(rendered_part).name
+        return rendered_part
+
+    def _render_parts(
+        self,
+        parts: tuple[str, ...],
+        rendered_parts: tuple[str, ...] | None = None,
+        extra_context: AnyByStrDict | None = None,
+        is_template: bool = False,
+    ) -> Iterable[tuple[Path, AnyByStrDict | None]]:
+        """Render a set of parts into path and context pairs.
+
+        If a yield tag is found in a part, it will recursively yield multiple path and context pairs.
+        """
+        if rendered_parts is None:
+            rendered_parts = tuple()
+
+        if not parts:
+            rendered_path = Path(*rendered_parts)
+
+            templated_sibling = (
+                self.template.local_abspath
+                / f"{rendered_path}{self.template.templates_suffix}"
+            )
+            if is_template or not templated_sibling.exists():
+                yield rendered_path, extra_context
+
+            return
+
+        part = parts[0]
+        parts = parts[1:]
+
+        if not extra_context:
+            extra_context = {}
+
+        # If the `part` has a yield tag, `self.jinja_env` will be set with the yield name and iterable
+        rendered_part = self._render_string(part, extra_context=extra_context)
+
+        yield_name = self.jinja_env.yield_name
+        if yield_name:
+            for value in self.jinja_env.yield_iterable or ():
+                new_context = {**extra_context, yield_name: value}
+                rendered_part = self._render_string(part, extra_context=new_context)
+                rendered_part = self._adjust_rendered_part(rendered_part)
+
+                # Skip if any part is rendered as an empty string
+                if not rendered_part:
+                    continue
+
+                yield from self._render_parts(
+                    parts, rendered_parts + (rendered_part,), new_context, is_template
+                )
+
+            return
+
+        # Skip if any part is rendered as an empty string
+        if not rendered_part:
+            return
+
+        rendered_part = self._adjust_rendered_part(rendered_part)
+
+        yield from self._render_parts(
+            parts, rendered_parts + (rendered_part,), extra_context, is_template
+        )
+
+    def _render_path(self, relpath: Path) -> Iterable[tuple[Path, AnyByStrDict | None]]:
+        """Render one relative path into multiple path and context pairs.
 
         Args:
             relpath:
@@ -729,29 +818,11 @@ def _render_path(self, relpath: Path) -> Path | None:
         )
         # With an empty suffix, the templated sibling always exists.
         if templated_sibling.exists() and self.template.templates_suffix:
-            return None
+            return
         if self.template.templates_suffix and is_template:
             relpath = relpath.with_suffix("")
-        rendered_parts = []
-        for part in relpath.parts:
-            # Skip folder if any part is rendered as an empty string
-            part = self._render_string(part)
-            if not part:
-                return None
-            # {{ _copier_conf.answers_file }} becomes the full path; in that case,
-            # restore part to be just the end leaf
-            if str(self.answers_relpath) == part:
-                part = Path(part).name
-            rendered_parts.append(part)
-        result = Path(*rendered_parts)
-        if not is_template:
-            templated_sibling = (
-                self.template.local_abspath
-                / f"{result}{self.template.templates_suffix}"
-            )
-            if templated_sibling.exists():
-                return None
-        return result
+
+        yield from self._render_parts(relpath.parts, is_template=is_template)
 
     def _render_string(
         self, string: str, extra_context: AnyByStrDict | None = None

docs/comparisons.md

@@ -18,6 +18,7 @@ docs! We don't want to be biased, but it's easy that we tend to be:
 | Feature                                  | Copier                           | Cookiecutter                    | Yeoman        |
 | ---------------------------------------- | -------------------------------- | ------------------------------- | ------------- |
 | Can template file names                  | Yes                              | Yes                             | Yes           |
+| Can generate file structures in loops    | Yes                              | No                              | No            |
 | Configuration                            | Single YAML file[^1]             | Single JSON file                | JS module     |
 | Migrations                               | Yes                              | No                              | No            |
 | Programmed in                            | Python                           | Python                          | NodeJS        |