Transform logging statements

logging_refactoring.py

"""Transform logging statements.

Dependencies:
1. refactor: https://pypi.org/project/refactor/
"""

import ast
import re

import refactor
from refactor.context import Scope


class ConvertFStringsToStringFormat(refactor.Rule):
    """Replace f-string messages in logging statement with format string based messages."""

    def match(self, node: ast.AST) -> refactor.BaseAction:
        assert isinstance(node, ast.Call)
        assert isinstance(node.func, ast.Attribute)
        assert node.func.attr in ("error", "info", "debug", "warning", "exception")
        assert isinstance(node.func.value, ast.Name)
        assert node.func.value.id.endswith("logger")
        assert node.args
        msg = node.args[0]
        assert isinstance(msg, ast.JoinedStr)
        fmt_str = "".join(
            (f.value if isinstance(f, ast.Constant) else "%s") for f in msg.values
        )
        values = [f.value for f in msg.values if isinstance(f, ast.FormattedValue)]
        result = ast.Call(
            func=node.func,
            args=[ast.Constant(fmt_str), *values],
            keywords=node.keywords,
        )
        return refactor.Replace(node, result)


class ConvertErrorWithExcInfoToException(refactor.Rule):
    """Replace logger.error(...,exc_info=var|True) with logger.exception(...)."""

    def match(self, node: ast.AST) -> refactor.BaseAction:
        assert isinstance(node, ast.Call)
        assert isinstance(node.func, ast.Attribute)
        assert node.func.attr == "error"
        assert node.keywords
        assert (kw := next((k for k in node.keywords if k.arg == "exc_info"), None))
        assert not isinstance(kw.value, ast.Constant) or kw.value.value
        result = ast.Call(
            func=ast.Attribute(value=node.func.value, attr="exception"),
            args=node.args,
            keywords=[k for k in node.keywords if k.arg != "exc_info"],
        )
        return refactor.Replace(node, result)


class Defs(refactor.Representative):
    context_providers = (Scope,)

    def is_exception(self, name: str, use_node: ast.AST) -> bool:
        defs = self.context.scope.resolve(use_node).get_definitions(name)
        return any(
            isinstance(node, ast.ExceptHandler) and node.name == name and node in defs
            for node in ast.walk(self.context.tree)
        )


class ConvertErrorWithExceptionAsLastArgumentToException(refactor.Rule):
    """Replace logger.error(...,ex) with logger.exception(...)."""

    context_providers = (Defs,)

    def match(self, node: ast.AST) -> refactor.BaseAction:
        assert isinstance(node, ast.Call)
        assert isinstance(node.func, ast.Attribute)
        assert node.func.attr == "error"
        assert len(node.args) > 1
        last_arg = node.args[-1]
        assert isinstance(last_arg, ast.Name)
        assert self.context.defs.is_exception(last_arg.id, node)
        arg1 = node.args[0]
        assert isinstance(arg1, ast.Constant)
        result = ast.Call(
            func=ast.Attribute(value=node.func.value, attr="exception"),
            args=[ast.Constant(re.sub(r":? ?%s$", "", arg1.value))] + node.args[1:-1],
            keywords=node.keywords,
        )
        return refactor.Replace(node, result)


class RemoveUnusedExceptionVariable(refactor.Rule):
    """Conservatively remove unused exception variables."""

    class Uses(refactor.Representative):
        context_providers = (Scope,)

        def is_used(self, name: str, def_node: ast.AST) -> bool:
            return any(
                isinstance(node, ast.Name) and node.id == name
                for node in ast.walk(def_node)
            )

    context_providers = (Uses,)

    def match(self, node: ast.AST) -> refactor.BaseAction:
        assert isinstance(node, ast.ExceptHandler)
        assert node.name
        assert not self.context.uses.is_used(node.name, node)
        result = ast.ExceptHandler(type=node.type, name=None, body=node.body)
        return refactor.Replace(node, result)


class ConvertMultilineMsgIntoSingleLineMsg(refactor.Rule):
    """Convert multiline message in logging statement to single line message."""

    def match(self, node: ast.AST) -> refactor.BaseAction:
        assert isinstance(node, ast.Call)
        assert isinstance(node.func, ast.Attribute)
        assert node.func.attr in ("error", "info", "debug", "warning", "exception")
        assert isinstance(node.func.value, ast.Name)
        assert node.func.value.id.endswith("logger")
        assert node.args
        msg = node.args[0]
        assert isinstance(msg, ast.Constant)
        new_msg = re.sub(r"\n +", " ", msg.value)
        assert msg.value != new_msg
        result = ast.Call(
            func=node.func,
            args=[ast.Constant(new_msg), *node.args[1:]],
            keywords=node.keywords,
        )
        return refactor.Replace(node, result)


class ConvertErrorWithExceptionInMessageToException(refactor.Rule):
    """Replace logger.error(msg,..,ex,...) with logger.exception(...)."""

    context_providers = (Defs,)

    def match(self, node: ast.AST) -> refactor.BaseAction:
        assert isinstance(node, ast.Call)
        assert isinstance(node.func, ast.Attribute)
        assert node.func.attr == "error"
        assert len(node.args) > 1
        msg = node.args[0]
        assert isinstance(msg, ast.Constant)
        remaining_args = [
            a
            for a in node.args[1:]
            if not (
                isinstance(a, ast.Name) and self.context.defs.is_exception(a.id, node)
            )
        ]
        assert len(remaining_args) == len(node.args) - 2
        ex_in_msg = r" ?\| Ex(ception)? ?(Traceback)?: %s ?"
        new_msg = re.sub(ex_in_msg, " ", msg.value)
        result = ast.Call(
            func=ast.Attribute(value=node.func.value, attr="exception"),
            args=[ast.Constant(new_msg)] + remaining_args,
            keywords=node.keywords,
        )
        return refactor.Replace(node, result)


if __name__ == "__main__":
    refactor.run(
        rules=[
            ConvertFStringsToStringFormat,
            ConvertErrorWithExcInfoToException,
            ConvertErrorWithExceptionAsLastArgumentToException,
            ConvertErrorWithExceptionInMessageToException,
            RemoveUnusedExceptionVariable,
            ConvertMultilineMsgIntoSingleLineMsg,
        ]
    )