bikini-patchwork/patchwork/transforms/identifiers.py

from __future__ import annotations
import ast
import builtins
import keyword
import random
from ..util import gen_name
_DUNDERS = {'__name__', '__main__', '__file__', '__doc__', '__builtins__', '__package__', '__path__', '__loader__', '__spec__', '__cached__', '__all__', '__init__', '__new__', '__del__', '__init_subclass__', '__subclasshook__', '__set_name__', '__class_getitem__', '__enter__', '__exit__', '__aenter__', '__aexit__', '__await__', '__aiter__', '__anext__', '__call__', '__repr__', '__str__', '__bytes__', '__format__', '__hash__', '__bool__', '__len__', '__length_hint__', '__sizeof__', '__dir__', '__eq__', '__ne__', '__lt__', '__le__', '__gt__', '__ge__', '__add__', '__sub__', '__mul__', '__truediv__', '__floordiv__', '__mod__', '__divmod__', '__pow__', '__lshift__', '__rshift__', '__and__', '__or__', '__xor__', '__radd__', '__rsub__', '__rmul__', '__rtruediv__', '__rfloordiv__', '__rmod__', '__rdivmod__', '__rpow__', '__rlshift__', '__rrshift__', '__rand__', '__ror__', '__rxor__', '__iadd__', '__isub__', '__imul__', '__itruediv__', '__ifloordiv__', '__imod__', '__ipow__', '__ilshift__', '__irshift__', '__iand__', '__ior__', '__ixor__', '__matmul__', '__rmatmul__', '__imatmul__', '__neg__', '__pos__', '__abs__', '__invert__', '__int__', '__float__', '__complex__', '__round__', '__index__', '__trunc__', '__floor__', '__ceil__', '__getattr__', '__setattr__', '__delattr__', '__getattribute__', '__getitem__', '__setitem__', '__delitem__', '__missing__', '__iter__', '__next__', '__reversed__', '__contains__', '__copy__', '__deepcopy__', '__reduce__', '__reduce_ex__', '__getstate__', '__setstate__', '__getnewargs__', '__getnewargs_ex__', '__class__', '__instancecheck__', '__subclasscheck__', '__slots__', '__dict__', '__weakref__', '__module__', '__qualname__', '__defaults__', '__kwdefaults__', '__annotations__', '__wrapped__', '__signature__', '__match_args__', '__post_init__', '__fspath__', '__index__', '__buffer__', '__release_buffer__'}
RESERVED: set[str] = set(dir(builtins)) | set(keyword.kwlist) | _DUNDERS | {'self', 'cls', 'metaclass'}

def _has_future_annotations(tree: ast.AST) -> bool:
    if not isinstance(tree, ast.Module):
        return False
    for stmt in tree.body:
        if isinstance(stmt, ast.Expr) and isinstance(stmt.value, ast.Constant) and isinstance(stmt.value.value, str):
            continue
        if isinstance(stmt, ast.ImportFrom) and stmt.module == '__future__':
            if any(alias.name == 'annotations' for alias in stmt.names):
                return True
            continue
        break
    return False

def _names_in_annotation(node: ast.AST | None) -> set[str]:
    if node is None:
        return set()
    return {child.id for child in ast.walk(node) if isinstance(child, ast.Name)}

def collect_skip_names(tree: ast.AST) -> set[str]:
    skip: set[str] = set()
    preserve_annotation_names = _has_future_annotations(tree)
    for node in ast.walk(tree):
        if isinstance(node, ast.Assign):
            for target in node.targets:
                if isinstance(target, ast.Name) and target.id == '__all__':
                    if isinstance(node.value, (ast.List, ast.Tuple, ast.Set)):
                        for elt in node.value.elts:
                            if isinstance(elt, ast.Constant) and isinstance(elt.value, str):
                                skip.add(elt.value)
        elif isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and (node.func.id in ('getattr', 'setattr', 'hasattr', 'delattr')):
            if len(node.args) >= 2:
                arg = node.args[1]
                if isinstance(arg, ast.Constant) and isinstance(arg.value, str):
                    skip.add(arg.value)
        elif isinstance(node, ast.ClassDef):
            for stmt in node.body:
                if isinstance(stmt, (ast.FunctionDef, ast.AsyncFunctionDef)):
                    skip.add(stmt.name)
                elif isinstance(stmt, ast.Assign):
                    for tgt in stmt.targets:
                        if isinstance(tgt, ast.Name):
                            skip.add(tgt.id)
                        elif isinstance(tgt, (ast.Tuple, ast.List)):
                            for elt in tgt.elts:
                                if isinstance(elt, ast.Name):
                                    skip.add(elt.id)
                elif isinstance(stmt, ast.AnnAssign) and isinstance(stmt.target, ast.Name):
                    skip.add(stmt.target.id)
        elif isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.Lambda)):
            args = node.args
            if preserve_annotation_names:
                if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
                    skip.update(_names_in_annotation(node.returns))
                for a in args.args + args.posonlyargs + args.kwonlyargs:
                    skip.update(_names_in_annotation(a.annotation))
                if args.vararg:
                    skip.update(_names_in_annotation(args.vararg.annotation))
                if args.kwarg:
                    skip.update(_names_in_annotation(args.kwarg.annotation))
            for a in args.args + args.posonlyargs + args.kwonlyargs:
                skip.add(a.arg)
            if args.vararg:
                skip.add(args.vararg.arg)
            if args.kwarg:
                skip.add(args.kwarg.arg)
        elif preserve_annotation_names and isinstance(node, ast.AnnAssign):
            skip.update(_names_in_annotation(node.annotation))
    return skip

class IdentifierRenamer(ast.NodeTransformer):

    def __init__(self, rng: random.Random, keep: set[str] | None=None):
        self.rng = rng
        self.mapping: dict[str, str] = {}
        self.reserved = RESERVED | (keep or set())
        self._used_outputs: set[str] = set()

    def _new(self) -> str:
        n = gen_name(self.rng, used=self._used_outputs)
        self._used_outputs.add(n)
        return n

    def _rename(self, name: str) -> str:
        if not name or name in self.reserved:
            return name
        if name.startswith('__') and name.endswith('__') and (len(name) >= 4):
            return name
        if name not in self.mapping:
            self.mapping[name] = self._new()
        return self.mapping[name]

    def visit_Name(self, node: ast.Name) -> ast.AST:
        node.id = self._rename(node.id)
        return node

    def visit_FunctionDef(self, node: ast.FunctionDef) -> ast.AST:
        node.name = self._rename(node.name)
        node.decorator_list = [self.visit(d) for d in node.decorator_list]
        if node.returns is not None:
            node.returns = self.visit(node.returns)
        node.args = self.visit(node.args)
        node.body = [self.visit(stmt) for stmt in node.body]
        return node
    visit_AsyncFunctionDef = visit_FunctionDef

    def visit_ClassDef(self, node: ast.ClassDef) -> ast.AST:
        node.name = self._rename(node.name)
        node.bases = [self.visit(b) for b in node.bases]
        for kw in node.keywords:
            kw.value = self.visit(kw.value)
        node.decorator_list = [self.visit(d) for d in node.decorator_list]
        node.body = [self.visit(stmt) for stmt in node.body]
        return node

    def visit_Lambda(self, node: ast.Lambda) -> ast.AST:
        node.args = self.visit(node.args)
        node.body = self.visit(node.body)
        return node

    def visit_arg(self, node: ast.arg) -> ast.AST:
        node.arg = self._rename(node.arg)
        if node.annotation is not None:
            node.annotation = self.visit(node.annotation)
        return node

    def visit_arguments(self, node: ast.arguments) -> ast.AST:
        node.args = [self.visit(a) for a in node.args]
        node.posonlyargs = [self.visit(a) for a in node.posonlyargs]
        node.kwonlyargs = [self.visit(a) for a in node.kwonlyargs]
        if node.vararg is not None:
            node.vararg = self.visit(node.vararg)
        if node.kwarg is not None:
            node.kwarg = self.visit(node.kwarg)
        node.defaults = [self.visit(d) for d in node.defaults]
        node.kw_defaults = [self.visit(d) if d is not None else None for d in node.kw_defaults]
        return node

    def visit_Attribute(self, node: ast.Attribute) -> ast.AST:
        node.value = self.visit(node.value)
        return node

    def visit_keyword(self, node: ast.keyword) -> ast.AST:
        if node.value is not None:
            node.value = self.visit(node.value)
        return node

    def visit_Global(self, node: ast.Global) -> ast.AST:
        node.names = [self._rename(n) for n in node.names]
        return node

    def visit_Nonlocal(self, node: ast.Nonlocal) -> ast.AST:
        node.names = [self._rename(n) for n in node.names]
        return node

    def visit_ExceptHandler(self, node: ast.ExceptHandler) -> ast.AST:
        if node.type is not None:
            node.type = self.visit(node.type)
        if node.name:
            node.name = self._rename(node.name)
        node.body = [self.visit(s) for s in node.body]
        return node

    def visit_Import(self, node: ast.Import) -> ast.AST:
        for alias in node.names:
            if '.' in alias.name:
                if alias.asname:
                    new = self._rename(alias.asname)
                    alias.asname = new
                continue
            if alias.asname:
                alias.asname = self._rename(alias.asname)
            else:
                if alias.name in self.reserved:
                    continue
                if alias.name not in self.mapping:
                    self.mapping[alias.name] = self._new()
                alias.asname = self.mapping[alias.name]
        return node

    def visit_ImportFrom(self, node: ast.ImportFrom) -> ast.AST:
        if node.module == '__future__':
            return node
        for alias in node.names:
            if alias.name == '*':
                continue
            if alias.asname:
                alias.asname = self._rename(alias.asname)
            else:
                if alias.name in self.reserved:
                    continue
                if alias.name not in self.mapping:
                    self.mapping[alias.name] = self._new()
                alias.asname = self.mapping[alias.name]
        return node

    def visit_MatchAs(self, node: ast.MatchAs) -> ast.AST:
        if node.pattern is not None:
            node.pattern = self.visit(node.pattern)
        if node.name is not None:
            node.name = self._rename(node.name)
        return node

    def visit_MatchStar(self, node: ast.MatchStar) -> ast.AST:
        if node.name is not None:
            node.name = self._rename(node.name)
        return node

    def visit_MatchMapping(self, node: ast.MatchMapping) -> ast.AST:
        node.keys = [self.visit(key) for key in node.keys]
        node.patterns = [self.visit(pattern) for pattern in node.patterns]
        if node.rest is not None:
            node.rest = self._rename(node.rest)
        return node