Check for usage of global variables in module functions

purity.py

"""
Check for usage of global variables in module functions.

WARNING: if a global variable is created inside `if __name__ == '__main__'` block
         it is not detected
"""
import ast
import inspect
import pkgutil
import sys
from contextlib import redirect_stderr
from importlib import import_module, util
from logging import Logger
from pathlib import Path
from types import ModuleType

SHOW_PURE = True


def bold(text: str) -> str:
    "Make text bold."
    return f"\033[1m{text}\033[0m"


def check_module_purity(module: ModuleType, *, show_pure: bool = True) -> None:
    """
    Check for usage of global variables in module functions.

    Skip functions, modules, classes.
    """
    functions = [  # Retrieve functions in the specified module
        (name, obj) for name, obj in inspect.getmembers(module)
        if inspect.isfunction(obj) and obj.__module__ == module.__name__
    ]

    for func_name, func_obj in functions:
        # Retrieve names of external variables which are used inside a function
        external_names = func_obj.__code__.co_names
        illegal_globals = []
        consts = []

        for name in external_names:
            obj = getattr(module, name, None)  # Get object by its name
            if obj is not None:
                # Filter out functions, modules, classes, consts, loggers
                is_func = inspect.isroutine(obj)  # func/method
                is_module = inspect.ismodule(obj)
                is_class = inspect.isclass(obj)
                is_const = name.isupper()
                is_logger = isinstance(obj, Logger)

                if not (is_func or is_module or is_class or is_const or is_logger):
                    illegal_globals.append(name)
                if is_const:
                    consts.append(name)

        str_consts = (' | consts: ' + ', '.join(consts)) if consts else ''
        if illegal_globals:
            print(f"❌ {bold(func_name)} | global vars: "
                  f"{', '.join(illegal_globals)}{str_consts}")
        elif show_pure:
            print(f"✅ {bold(func_name)}{str_consts}")


if __name__ == "__main__":
    if Path(sys.argv[-1]).is_file():
        spec = util.spec_from_file_location("user_module", sys.argv[-1])
        if not spec:
            msg = "cannot initialize spec"
            raise ValueError(msg)
        mod = util.module_from_spec(spec)
        # sys.modules["user_module"] = mod
        if spec.loader:
            spec.loader.exec_module(mod)
    else:
        mod = import_module(sys.argv[-1])

    if not hasattr(mod, "__path__"):
        print(f"Test module {bold(mod.__name__)}")
        check_module_purity(mod, show_pure=SHOW_PURE)
    else:
        mods = []
        for i in pkgutil.walk_packages(mod.__path__, mod.__name__ + "."):
            print(f"Test module {bold(i.name)}")
            try:
                with redirect_stderr(None):
                    check_module_purity(import_module(i.name), show_pure=SHOW_PURE)
            except Exception as e:  # noqa: BLE001
                print("Failed to import:", e)
            print()