jinja_unused_variable_finder.py

"""
Copyright (c) 2021 [Son Geon <work.jen6@gmail.com>]

Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
"""

from dataclasses import dataclass
from unittest.case import TestCase

from jinja2.meta import TrackingCodeGenerator
from jinja2 import Environment, Template, DictLoader
from jinja2.nodes import Name
from collections import namedtuple

TemplateVariables = namedtuple("TemplateVariables", "managed unmanaged")


class VariableFinder(TrackingCodeGenerator):
    def __init__(self, environment):
        super().__init__(environment)
        self.import_namespaces = set()
        self.macros = set()
        self.contexts = []
        self.context_managed_identifiers = set()

    def write(self, x):
        """Don't write."""

    def visit_Import(self, node, frame):
        super().visit_Import(node, frame)
        self.import_namespaces.add(node.target)

    def visit_Assign(self, node, frame):
        super().visit_Assign(node, frame)
        self.import_namespaces.add(node.target.name)

    def visit_Macro(self, node, frame):
        # macro definition에서 첫번째 인자의 이름을 context로 가져옵니다.
        if (
            len(node.args) != 0
            and node.args[0].ctx == "param"
            and node.args[0].name != ""
        ):
            """
            macro에서 받는 파라미터가 존재할 경우 안에서 사용하는 변수를 managed context로 인식하기 위해 Macro 내부의 ASt
            """
            self.contexts.append(node.args[0].name)
            super().visit_Macro(node, frame)
            self.contexts.pop()
        else:
            super().visit_Macro(node, frame)

    def visit_Getattr(self, node, frame):
        attrs = []
        last_node = node
        while True:
            if isinstance(last_node, Name):
                attrs.append(last_node.name)
                break
            attrs.append(last_node.attr)
            last_node = last_node.node

        attrs.reverse()
        attr_name = ".".join(attrs)

        if last_node.name in self.import_namespaces:
            return

        # macro context안에서 관리되는 변수인 경우 context_managed_identifiers 에 넣어줍니다.
        if len(self.contexts) != 0 and last_node.name == self.contexts[-1]:
            self.context_managed_identifiers.add(attr_name)
        else:
            self.undeclared_identifiers.add(attr_name)

    def pull_locals(self, frame):
        # import namespace는 undeclare에서 제외시켜준다
        super().pull_locals(frame)
        for namespace in self.import_namespaces:
            if namespace in self.undeclared_identifiers:
                self.undeclared_identifiers.remove(namespace)

    def get_undeclared_variables(self, ast) -> TemplateVariables:
        self.visit(ast)
        for namespace in self.import_namespaces:
            if namespace in self.undeclared_identifiers:
                self.undeclared_identifiers.remove(namespace)
        result = TemplateVariables(
            self.context_managed_identifiers, self.undeclared_identifiers
        )
        return result


class TestTemplateRender(TestCase):
    def get_undeclared_variables(
        self, env: Environment, template_name: str
    ) -> TemplateVariables:
        source = env.loader.get_source(env, template_name)[0]
        ast = env.parse(source)
        variable_finder = VariableFinder(ast.environment)
        return variable_finder.get_undeclared_variables(ast)

    def get_unused_variable_in_context(self, used_keys, context):
        context_keys = set(context.keys())
        for key in used_keys:
            context_keys.remove(key)
        return context_keys

    def check_variables_in_style_context(
        self, undeclared_variable, style_context
    ) -> bool:
        keys = undeclared_variable.split(".")
        current_context = style_context
        for key in keys:
            if isinstance(current_context, dict) and key in current_context:
                current_context = current_context[key]
            elif hasattr(current_context, key):
                current_context = getattr(current_context, key)
            else:
                return False
        return True

    def check_template_with_context(
        self, undeclared_variables, context, template_name: str, check_managed=False
    ):
        """
        template에서 사용되는 undeclared_variable들이 context내에 선언되지 않은 경우를 검사합니다.
        또한 check_managed=True 일 때 managed 변수만 사용하는지와 context내에 template에서 사용하지
        않는 경우가 있는지를 확인합니다.
        """
        used_key = set()
        for undeclared_variable in undeclared_variables.managed:
            if self.check_variables_in_style_context(undeclared_variable, context):
                key = undeclared_variable.split(".")[0]
                used_key.add(key)
            else:
                self.fail(
                    f"{template_name} has undefined variables : {undeclared_variable}"
                )

        for undeclared_variable in undeclared_variables.unmanaged:
            if self.check_variables_in_style_context(undeclared_variable, context):
                key = undeclared_variable.split(".")[0]
                used_key.add(key)
            else:
                self.fail(
                    f"{template_name} has undefined variables : {undeclared_variable}"
                )

        if check_managed:
            unused_vairables = self.get_unused_variable_in_context(used_key, context)
            self.assertEqual(
                len(unused_vairables),
                0,
                f"{template_name} has unused variables {unused_vairables}\n used variables : {used_key}",
            )

            if len(undeclared_variables.unmanaged) != 0:
                self.fail(
                    f"{template_name} has unmanaged variables : {undeclared_variables.unmanaged}"
                )

    def test_template_undeclared_context(self):
        main_template = """
        <h1> {{app_subdomain}} </h1>
        """

        loader = DictLoader({"template": main_template,})
        env = Environment(loader=loader)
        undeclared_variables = self.get_undeclared_variables(env, "template")

        # pass
        self.check_template_with_context(
            undeclared_variables=undeclared_variables,
            context={"app_subdomain": "subdomain"},
            template_name="template",
            check_managed=False,
        )

        # AssertionError: template has undefined variables :  app_subdomain
        self.check_template_with_context(
            undeclared_variables=undeclared_variables,
            context={},
            template_name="template",
            check_managed=False,
        )

    def test_template_unused_variable_exists_in_context(self):
        main_template = """
        <h1> {{app_subdomain}} </h1>
        """

        loader = DictLoader({"template": main_template,})
        env = Environment(loader=loader)
        undeclared_variables = self.get_undeclared_variables(env, "template")

        # pass
        self.check_template_with_context(
            undeclared_variables=undeclared_variables,
            context={"app_subdomain": "subdomain"},
            template_name="template",
            check_managed=False,
        )

        # AssertionError: 1 != 0 : template has unused variables {'unused_var'}
        self.check_template_with_context(
            undeclared_variables=undeclared_variables,
            context={"app_subdomain": "subdomain", "unused_var": "unused"},
            template_name="template",
            check_managed=True,
        )

    def test_template_unmanaged_context_exists(self):
        @dataclass
        class SdkToken:
            app_subdomain: str
            app_token: str

        managed_context_sdk = """
        {% macro sdkInit(sdk_context) -%}
          airbridge.init({
            'app': '{{sdk_context.app_subdomain}}',
            'token': '{{sdk_context.app_token}}',
          })
        {%- endmacro %}
        """

        unmanaged_context_sdk = """
        {% macro sdkInit(sdk_context) -%}
          airbridge.init({
            'app': '{{app_subdomain}}',
            'token': '{{sdk_context.app_token}}',
          })
        {%- endmacro %}
        """

        loader = DictLoader(
            {
                "managed_context_sdk": managed_context_sdk,
                "unmanaged_context_sdk": unmanaged_context_sdk,
            }
        )
        env = Environment(loader=loader)

        # pass
        context = {
            "sdk_context": SdkToken(app_subdomain="subdomain", app_token="token"),
        }
        undeclared_variables = self.get_undeclared_variables(env, "managed_context_sdk")
        self.check_template_with_context(
            undeclared_variables, context, "managed_context_sdk", check_managed=True
        )

        # AssertionError: unmanaged_context_sdk has unmanaged variables : {'app_subdomain'}
        context = {
            "app_subdomain": "subdomain",
            "sdk_context": SdkToken(app_subdomain="subdomain", app_token="token"),
        }
        undeclared_variables = self.get_undeclared_variables(
            env, "unmanaged_context_sdk"
        )
        self.check_template_with_context(
            undeclared_variables, context, "unmanaged_context_sdk", check_managed=True
        )