from functools import wraps

from dbt.plugins.manager import dbtPlugin
from dbt_common.clients import jinja as dbt_jinja
from jinja2 import StrictUndefined, Undefined


class DbtStrictUndefinedPlugin(dbtPlugin):
    """Adds Python's subprocess module to dbt's `modules` context."""

    def initialize(self) -> None:
        self._wrap_dbt_jinja_environment_factory()

    def _wrap_dbt_jinja_environment_factory(self):
        self._orig_get_environment = dbt_jinja.get_environment

        @wraps(self._orig_get_environment)
        def _get_environment_with_strict_undefined(*a, **kw):
            jinja_env = self._orig_get_environment(*a, **kw)

            # only substitute if there's no custom logic in there already
            # (e.g. for macro tracing at parse time)
            if jinja_env.undefined is Undefined:
                jinja_env.undefined = StrictUndefined

            return jinja_env

        dbt_jinja.get_environment = _get_environment_with_strict_undefined


plugins = [DbtStrictUndefinedPlugin]
