Source code for codegen.env

"""Jinja2 environment creation and template rendering."""

from typing import TYPE_CHECKING

import jinja2

if TYPE_CHECKING:
    from collections.abc import Mapping
    from pathlib import Path

    from jinja2.runtime import Context

    from codegen.imports import ImportCollector
    from codegen.naming import Name




def create_jinja_env(
    *template_dirs: Path,
    extra_globals: Mapping[str, object] | None = None,
) -> jinja2.Environment:
    env = jinja2.Environment(
        loader=jinja2.FileSystemLoader(
            [str(template_dir) for template_dir in template_dirs],
        ),
        extensions=["jinja2.ext.do"],
        trim_blocks=True,
        lstrip_blocks=True,
        keep_trailing_newline=True,
        undefined=jinja2.StrictUndefined,
        autoescape=False,  # noqa: S701 — generating source, not HTML
    )
    env.globals["import_from"] = _import_from  # type: ignore[assignment]
    env.globals["import_name"] = _import_name  # type: ignore[assignment]

    if extra_globals:
        env.globals.update(extra_globals)  # type: ignore[arg-type]

    return env



def _collector(ctx: Context) -> ImportCollector:
    collector = ctx.get("_imports")

    if collector is None:
        msg = (
            "import_from/import_module called with no active import "
            "collector; add `with context` to the `{% import %}` / "
            "`{% from ... import %}` of the macro that declares imports"
        )
        raise RuntimeError(msg)

    return collector


@jinja2.pass_context
def _import_from(ctx: Context, module: str, *names: str) -> str:
    _collector(ctx).add_from(module, *names)
    return ""


@jinja2.pass_context
def _import_name(ctx: Context, *names: Name) -> str:
    _collector(ctx).add_name(*names)
    return ""




def render_template(
    env: jinja2.Environment,
    template_name: str,
    **context: object,
) -> str:
    try:
        return env.get_template(template_name).render(**context)

    except Exception as exception:
        raise Exception(
            str(exception),
            {
                key: value
                for key, value in context.items()
                if key not in {"ctx", "_imports", "import_block"}
            },
        ) from exception