logging.* auto refactoring

logging_refactor.py

import ast
import copy
from pathlib import Path

import asttokens
import astunparse
import string

LOGGING_METHODS = {'debug', 'info', 'warn', 'warning', 'exception', 'error', 'fatal', 'critical'}


def clone(node):
    return copy.deepcopy(node)


class Visitor(ast.NodeVisitor):
    def __init__(self):
        self.transforms = []

    def visit_Call(self, node):
        if isinstance(node.func, ast.Attribute) and isinstance(node.func.value, ast.Name) and node.func.value.id == 'logging' and node.func.attr in LOGGING_METHODS:
            node = self.fix_logging(node)
        return node

    def fix_logging(self, node):
        n = clone(node)
        transformed = False
        if node.func.attr == 'warn':
            # 'warn' is deprecated
            n.func.attr = 'warning'
            transformed = True
        if len(node.args) == 1:
            # suspicious
            arg = n.args[0]
            if isinstance(arg, ast.Call) and isinstance(arg.func, ast.Attribute) and isinstance(arg.func.value, ast.Str) and arg.func.attr == 'format':
                # 'foo {}'.format(whatever)
                n.args = self.fix_logging_format(arg)
                transformed = True
            elif isinstance(arg, ast.BinOp) and isinstance(arg.op, ast.Mod) and isinstance(arg.left, ast.Str):
                # 'foo %s' % (whatever)
                n.args = self.fix_logging_mod(arg)
                transformed = True
        if transformed:
            self.transforms.append((node, n))
        return node

    def fix_logging_format(self, arg):
        fmt = arg.func.value.s
        new_fmt = []
        if arg.keywords:
            raise ValueError(".format(**keywords) is not supported")
        for i, (text, name, spec, conversion) in enumerate(string.Formatter().parse(fmt)):
            new_fmt.append(text)
            if name is None:
                continue
            new_fmt.append("%")
            if conversion == "r":
                new_fmt.append("r")
            elif (spec or "").endswith("%"):
                # seriously tho.
                new_fmt.append(spec[:-1] + "f%%")
                arg.args[i] = ast.BinOp(left=arg.args[i], op=ast.Mult(), right=ast.Num(n=100))
            else:
                new_fmt.append(spec or "s")
        return [ast.Str(s=''.join(new_fmt))] + arg.args

    def fix_logging_mod(self, arg):
        if isinstance(arg.right, (ast.Tuple, ast.List, ast.Set)):
            args = arg.right.elts
        else:
            args = [arg.right]
        return [arg.left] + args


def refactor(path):
    source = path.open().read()
    atok = asttokens.ASTTokens(source, parse=True)
    trans = Visitor()
    trans.visit(atok.tree)
    if not trans.transforms:
        return False

    source = list(source)
    diff = 0
    for prev, new in trans.transforms:
        new_text = astunparse.unparse(new).rstrip("\n")
        start, end = atok.get_text_range(prev)
        start += diff
        end += diff
        source[start:end] = new_text
        diff += len(new_text) - (end - start)

    with path.open('w') as f:
        f.write("".join(source))

    return True

if __name__ == '__main__':
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument('dir', nargs='+')
    args = parser.parse_args()

    for dir in args.dir:
        for path in Path(dir).rglob('*.py'):
            if refactor(path):
                print(path)