import ast

import copy

import keyword

from collections import Counter

from typing import Dict, List, Optional, Set

from .utils import variable_name_generator

class Transformer:

def transform(self, *trees):

for tree in trees:

self.visit(tree)

return trees

class NodeTransformer(Transformer, ast.NodeTransformer):

pass

class NodeVisitor(Transformer, ast.NodeVisitor):

pass

class Pipeline:

def __init__(self, *transformers):

self.transformers = transformers

def transform(self, *trees):

for transformer in self.transformers:

trees = transformer.transform(*trees)

return trees

class ReturnSimplifier(NodeTransformer):

"""Simplify return statements in the following form:

x = (some code)

return x

to

return (some code)

NOTE: unused_assignments must be modified in-place, since the set is passed

to RemoveUnusedVariables at initialization. Can't return a new set.

"""

def __init__(self, *args, **kwargs):

super().__init__(*args, **kwargs)

self.unused_assignments = set()

def _can_simplify_return(self, previous: ast.stmt, current: ast.stmt) -> bool:

return (

isinstance(previous, ast.Assign)

and len(previous.targets) == 1

and isinstance(previous.targets[0], ast.Name)

and isinstance(current, ast.Return)

and isinstance(current.value, ast.Name)

and current.value.id == previous.targets[0].id

)

def _simplify_body(self, body: List[ast.stmt]) -> List[ast.stmt]:

for previous, current in zip(body, body[1:]):

if self._can_simplify_return(previous, current):

self.unused_assignments.add(id(previous))

current.value = copy.deepcopy(previous.value)

return body

def generic_visit(self, node):

node = super().generic_visit(node)

for field, value in ast.iter_fields(node):

if isinstance(value, list) and value and all(isinstance(item, ast.stmt) for item in value):

setattr(node, field, self._simplify_body(value))

return node

class RemoveUnusedVariables(NodeTransformer):

"""Remove all unused variables.

NOTE: cannot store a copy of unused_assignments, as this set is modified

in-place

after initialization.

"""

def __init__(self, unused_assignments: Set[int]):

super().__init__()

self.unused_assignments = unused_assignments

def visit_Assign(self, node: ast.Assign) -> Optional[ast.Assign]:

if id(node) in self.unused_assignments:

return None

return self.generic_visit(node)

class VariableNameCollector(NodeVisitor):

"""Collects all variable names in scope."""

def __init__(self):

self.names = set()

def visit_Name(self, node):

self.names.add(node.id)

class ScopeLocalNameCollector(ast.NodeVisitor):

def __init__(self):

self.reserved_names = set()

self.bindings = set()

self.loads = set()

self.external_bindings = set()

self.args = set()

def visit_Name(self, node):

self.reserved_names.add(node.id)

if isinstance(node.ctx, ast.Store):

self.bindings.add(node.id)

elif isinstance(node.ctx, ast.Load):

self.loads.add(node.id)

def visit_arg(self, node):

self.reserved_names.add(node.arg)

self.bindings.add(node.arg)

self.args.add(node.arg)

if node.annotation is not None:

self.visit(node.annotation)

def visit_Global(self, node):

self.reserved_names.update(node.names)

self.external_bindings.update(node.names)

visit_Nonlocal = visit_Global

def visit_ExceptHandler(self, node):

if node.name:

self.reserved_names.add(node.name)

self.bindings.add(node.name)

if node.type is not None:

self.visit(node.type)

for statement in node.body:

self.visit(statement)

def visit_Import(self, node):

for alias in node.names:

bound_name = alias.asname or alias.name.split(".", 1)[0]

self.reserved_names.add(bound_name)

self.bindings.add(bound_name)

def visit_ImportFrom(self, node):

for alias in node.names:

if alias.name == "*":

continue

bound_name = alias.asname or alias.name

self.reserved_names.add(bound_name)

self.bindings.add(bound_name)

def _visit_nested_function(self, node):

self.reserved_names.add(node.name)

self.bindings.add(node.name)

for decorator in node.decorator_list:

self.visit(decorator)

for default in node.args.defaults:

self.visit(default)

for default in node.args.kw_defaults:

if default is not None:

self.visit(default)

for argument in (

[*node.args.posonlyargs, *node.args.args, *node.args.kwonlyargs]

+ ([node.args.vararg] if node.args.vararg is not None else [])

+ ([node.args.kwarg] if node.args.kwarg is not None else [])

):

if argument is not None and argument.annotation is not None:

self.visit(argument.annotation)

returns = getattr(node, "returns", None)

if returns is not None:

self.visit(returns)

def visit_FunctionDef(self, node):

self._visit_nested_function(node)

visit_AsyncFunctionDef = visit_FunctionDef

def visit_ClassDef(self, node):

self.reserved_names.add(node.name)

self.bindings.add(node.name)

for decorator in node.decorator_list:

self.visit(decorator)

for base in node.bases:

self.visit(base)

for keyword in node.keywords:

self.visit(keyword)

def visit_Lambda(self, node):

return None

def visit_ListComp(self, node):

return None

def visit_SetComp(self, node):

return None

def visit_DictComp(self, node):

return None

def visit_GeneratorExp(self, node):

return None

class ParentSetter(NodeTransformer):

"""Adds parent attribute to each node.

>>> def apply(src):

... tree = ast.parse(src)

... ParentSetter().visit(tree)

... return tree

...

>>> tree = apply("lorem = 'demiurgic'\\nipsum = 'demiurgic'")

>>> isinstance(tree.body[0].parent, ast.Module)

True

>>> isinstance(tree.body[0].value.parent, ast.Assign)

True

"""

def visit(self, node):

for child in ast.iter_child_nodes(node):

child.parent = node

self.visit(child)

return node

class CommentRemover(NodeTransformer):

"""Drop all comments, both single-line and docstrings.

>>> def apply(code):

... tree = ast.parse(code)

... tree = ParentSetter().visit(tree)

... tree = CommentRemover().visit(tree)

... return ast.unparse(tree)

...

>>> apply('1 + 1 # comment')

'1 + 1'

>>> print(apply('''

... def square(x):

... \\'\\'\\'Return the square of x.\\'\\'\\'

... return x ** 2

... '''))

def square(x):

return x ** 2

>>> print(apply('''

... def square(x):

... \\'\\'\\'Return the square of x.\\'\\'\\'

... '''))

def square(x):

0

"""

def visit_Expr(self, node):

if isinstance(node.value, ast.Constant) and isinstance(node.value.value, str):

if len(node.parent.body) == 1: # if body is just the comment

return ast.parse('0').body[0] # replace comment with 0

return None # otherwise, remove comment

return node

class VariableShortener(NodeTransformer):

"""Renames variables according to provided mapping.

>>> shortener = VariableShortener(variable_name_generator(), mapping={'donotrename': 'donotrename'})

>>> apply = lambda src: ast.unparse(shortener.visit(ast.parse(src)))

>>> print(apply('mamamia = 1; donotrename = 2;'))

a = 1

donotrename = 2

"""

# Compression passes in this transformer use these guardrails:

# - repeated-name aliasing is statement-local and deleted after the

# statement, so helpers do not leak across later code

# - repeated-string hoisting now runs at function, module, and class scope,

# with cleanup deletes for module/class helpers

# - preserve-public-API mode can rename top-level classes, methods, and

# class-body attributes, but it emits explicit aliases and fixes class

# __name__/__qualname__ for compatibility

# - attribute rewriting is limited to owners we can prove from the AST

# (`self`, `cls`, or known class names), not arbitrary dynamic receivers

#

# Keep regression coverage in tests/test_api.py and the checked-in example

# outputs under tests/examples in sync whenever these rules change.

def __init__(

self,

generator,

mapping=None,

modules=(),

keep_global_variables=False,

rename_arguments=False,

reserved_names=None,

):

self.mapping = mapping or {}

self.mapping_values = set(self.mapping.values())

self.generator = generator

self.reserved_names = set(reserved_names or ())

self.nodes_to_append = []

self.public_global_names = set()

self.scope_stack = []

self.local_rename_scopes = []

self.instance_type_scopes = [{}]

self.class_context_stack = []

self.class_member_mappings = {}

self.callable_argument_infos = {}

self.class_method_argument_infos = {}

self._class_public_member_reference_cache = {}

self._module_attribute_reference_cache = {}

self._scope_analysis_cache = {}

self.modules = set(modules) # don't alias variables imported from these modules

self.keep_global_variables = keep_global_variables

self.rename_arguments = rename_arguments

def _is_node_global(self, node):

"""Check if a node is global."""

return (

not hasattr(node, 'parent') or isinstance(node.parent, ast.Module)

)

def _rename_identifier(self, old_name):

if old_name not in self.mapping:

while True:

candidate = next(self.generator)

if candidate not in self.mapping_values and candidate not in self.reserved_names:

self.mapping[old_name] = candidate

self.mapping_values.add(candidate)

break

return self.mapping[old_name]

def _lookup_local_identifier(self, old_name):

for scope in reversed(self.local_rename_scopes):

if old_name in scope["mapping"]:

return scope["mapping"][old_name]

return None

def _lookup_visible_identifier(self, old_name):

local_name = self._lookup_local_identifier(old_name)

if local_name is not None:

return local_name

return self.mapping.get(old_name)

def _rename_local_identifier(self, old_name):

if old_name in self.mapping_values:

return old_name

scope = self.local_rename_scopes[-1]

if old_name not in scope["mapping"]:

new_name = next(scope["generator"])

scope["mapping"][old_name] = new_name

scope["used_names"].add(new_name)

return scope["mapping"][old_name]

def _push_instance_scope(self):

self.instance_type_scopes.append({})

def _pop_instance_scope(self):

self.instance_type_scopes.pop()

def _set_instance_type(self, name, class_name):

scope = self.instance_type_scopes[-1]

if class_name is None:

scope.pop(name, None)

else:

scope[name] = class_name

def _lookup_instance_type(self, name):

for scope in reversed(self.instance_type_scopes):

if name in scope:

return scope[name]

return None

def _append_public_alias(self, old_name, new_name):

if old_name != new_name:

self.nodes_to_append.append(ast.parse(f"{old_name} = {new_name}").body[0])

def _generated_assignment(self, source):

node = ast.parse(source).body[0]

node._pymini_generated = True

return node

def _containing_module(self, node):

current = node

while hasattr(current, "parent") and not isinstance(current.parent, ast.Module):

current = current.parent

return current.parent if hasattr(current, "parent") else None

def _current_class_context(self):

if self.class_context_stack:

return self.class_context_stack[-1]

return None

def _preserve_function_name(self, name):

return name.startswith("__") and name.endswith("__")

def _estimated_short_name_length(self):

return 1

def _rename_savings(self, old_name, count):

return max(0, len(old_name) - self._estimated_short_name_length()) * count

def _public_class_alias_cost(self, old_name):

short_name = "a"

return sum(

len(statement)

for statement in (

f"{old_name}={short_name}",

f"{short_name}.__name__={old_name!r}",

f"{short_name}.__qualname__={old_name!r}",

)

)

def _public_member_alias_cost(self, old_name):

return len(f"{old_name}=a")

def _public_global_alias_cost(self, old_name):

return len(f"{old_name}=a")

def _public_class_reference_count(self, node, old_name):

module = self._containing_module(node)

count = 1

if module is None:

return count

for current in ast.walk(module):

if isinstance(current, ast.Name) and current.id == old_name:

count += 1

return count

def _class_public_member_references(self, class_node):

cache_key = id(class_node)

cached = self._class_public_member_reference_cache.get(cache_key)

if cached is not None:

return cached

name_loads = Counter()

attribute_loads_by_base = {}

for current in ast.walk(class_node):

if isinstance(current, ast.Name) and isinstance(current.ctx, ast.Load):

name_loads[current.id] += 1

elif isinstance(current, ast.Attribute) and isinstance(current.value, ast.Name):

base_name = current.value.id

base_counts = attribute_loads_by_base.get(base_name)

if base_counts is None:

base_counts = Counter()

attribute_loads_by_base[base_name] = base_counts

base_counts[current.attr] += 1

cached = {

"name_loads": name_loads,

"attribute_loads_by_base": attribute_loads_by_base,

}

self._class_public_member_reference_cache[cache_key] = cached

return cached

def _module_attribute_references(self, module):

cache_key = id(module)

cached = self._module_attribute_reference_cache.get(cache_key)

if cached is not None:

return cached

attribute_loads_by_base = {}

for current in ast.walk(module):

if not isinstance(current, ast.Attribute) or not isinstance(current.value, ast.Name):

continue

base_name = current.value.id

base_counts = attribute_loads_by_base.get(base_name)

if base_counts is None:

base_counts = Counter()

attribute_loads_by_base[base_name] = base_counts

base_counts[current.attr] += 1

self._module_attribute_reference_cache[cache_key] = attribute_loads_by_base

return attribute_loads_by_base

def _public_member_reference_count(self, class_node, class_name, member_name):

references = self._class_public_member_references(class_node)

count = 1 + references["name_loads"].get(member_name, 0)

attribute_loads_by_base = references["attribute_loads_by_base"]

for base_name in {"self", "cls", class_name}:

count += attribute_loads_by_base.get(base_name, {}).get(member_name, 0)

module = self._containing_module(class_node)

if module is not None:

count += self._module_attribute_references(module).get(class_name, {}).get(member_name, 0)

return count

def _public_global_reference_count(self, node, old_name):

module = self._containing_module(node)

count = 1

if module is None:

return count

for current in ast.walk(module):

if (

isinstance(current, ast.Name)

and isinstance(current.ctx, ast.Load)

and current.id == old_name

):

count += 1

return count

def _should_rename_public_class(self, node, old_name):

return self._rename_savings(

old_name,

self._public_class_reference_count(node, old_name),

) > self._public_class_alias_cost(old_name)

def _should_rename_public_member(self, class_node, class_name, member_name):

return self._rename_savings(

member_name,

self._public_member_reference_count(class_node, class_name, member_name),

) > self._public_member_alias_cost(member_name)

def _should_rename_public_global(self, node, old_name):

return self._rename_savings(

old_name,

self._public_global_reference_count(node, old_name),

) > self._public_global_alias_cost(old_name)

def _is_method_definition(self, node):

return isinstance(getattr(node, "parent", None), ast.ClassDef)

def _is_class_body_assignment(self, node):

return isinstance(getattr(node, "parent", None), ast.ClassDef)

def _should_preserve_binding_targets(self, node):

return self.keep_global_variables and (

self._is_node_global(node) or self._is_class_body_assignment(node)

)

def _binding_names_from_target(self, target):

names = set()

if isinstance(target, ast.Name):

names.add(target.id)

elif isinstance(target, (ast.Tuple, ast.List)):

for element in target.elts:

names.update(self._binding_names_from_target(element))

return names

def _function_argument_nodes(self, arguments):

return [

*arguments.posonlyargs,

*arguments.args,

*arguments.kwonlyargs,

*([arguments.vararg] if arguments.vararg is not None else []),

*([arguments.kwarg] if arguments.kwarg is not None else []),

]

def _is_staticmethod(self, node):

for decorator in node.decorator_list:

if isinstance(decorator, ast.Name) and decorator.id == "staticmethod":

return True

if isinstance(decorator, ast.Attribute) and decorator.attr == "staticmethod":

return True

return False

def _should_rename_argument(self, name):

return (

self.rename_arguments

and not self._preserve_function_name(name)

and (len(name) > 1 or name in {"self", "cls"})

)

def _rename_function_arguments(self, node):

argument_mapping = {}

positional_params = [arg.arg for arg in [*node.args.posonlyargs, *node.args.args]]

for argument in self._function_argument_nodes(node.args):

old_name = argument.arg

if not self._should_rename_argument(old_name):

continue

new_name = self._rename_local_identifier(old_name)

argument.arg = new_name

if old_name != new_name:

argument_mapping[old_name] = new_name

receiver_names = set()

if self._is_method_definition(node) and not self._is_staticmethod(node) and positional_params:

receiver_name = positional_params.pop(0)

receiver_names.add(receiver_name)

receiver_names.add(argument_mapping.get(receiver_name, receiver_name))

return {

"rename_map": argument_mapping,

"positional_params": positional_params,

"receiver_names": receiver_names,

}

def _record_callable_argument_info(self, old_name, new_name, argument_info):

if not argument_info["rename_map"] and not argument_info["positional_params"]:

return

copied = {

"rename_map": dict(argument_info["rename_map"]),

"positional_params": list(argument_info["positional_params"]),

}

self.callable_argument_infos[old_name] = copied

self.callable_argument_infos[new_name] = copied

def _call_argument_info(self, func):

if isinstance(func, ast.Name):

return self.callable_argument_infos.get(func.id)

if not isinstance(func, ast.Attribute):

return None

base_name = func.value.id if isinstance(func.value, ast.Name) else None

class_context = self._current_class_context()

if base_name is not None:

if class_context is not None and base_name in (

class_context["receiver_names"]

| {class_context["old_name"], class_context["new_name"]}

):

return class_context["argument_infos"].get(func.attr)

instance_class = self._lookup_instance_type(base_name)

if instance_class in self.class_method_argument_infos:

return self.class_method_argument_infos[instance_class].get(func.attr)

if base_name in self.class_method_argument_infos:

return self.class_method_argument_infos[base_name].get(func.attr)

receiver_class = self._receiver_class_name(func.value)

if receiver_class in self.class_method_argument_infos:

return self.class_method_argument_infos[receiver_class].get(func.attr)

return None

def _rewrite_keywords_as_positional(self, node, argument_info):

if any(isinstance(arg, ast.Starred) for arg in node.args):

return

positional_params = argument_info["positional_params"]

if not positional_params:

return

next_position = len(node.args)

rewritten_keywords = []

can_convert = True

for keyword in node.keywords:

if (

can_convert

and keyword.arg is not None

and next_position < len(positional_params)

and keyword.arg == positional_params[next_position]

):

node.args.append(keyword.value)

next_position += 1

continue

can_convert = False

rewritten_keywords.append(keyword)

node.keywords = rewritten_keywords

def _rename_assignment_target(self, target, create_new=True):

if isinstance(target, ast.Name):

if self._is_active_parameter_name(target.id) or self._preserve_function_name(target.id):

return

if (

self.local_rename_scopes

and self.scope_stack

and target.id in self.scope_stack[-1]["bindings"]

and target.id not in self.scope_stack[-1]["globals"]

):

target.id = self._rename_local_identifier(target.id)

return

if target.id in self.mapping:

target.id = self.mapping[target.id]

elif create_new and target.id not in self.mapping_values:

target.id = self._rename_identifier(target.id)

return

if isinstance(target, (ast.Tuple, ast.List)):

for element in target.elts:

self._rename_assignment_target(element, create_new=create_new)

def _is_in_expression_scope(self, node):

current = getattr(node, "parent", None)

expression_scopes = (

ast.Lambda,

ast.ListComp,

ast.SetComp,

ast.DictComp,

ast.GeneratorExp,

)

while current is not None:

if isinstance(current, expression_scopes):

return True

current = getattr(current, "parent", None)

return False

def _is_in_function_signature(self, node):

current = getattr(node, "parent", None)

while current is not None:

if isinstance(current, (ast.FunctionDef, ast.AsyncFunctionDef)):

return False

if isinstance(current, (ast.arguments, ast.arg)):

return True

current = getattr(current, "parent", None)

return False

def _scope_analysis(self, node):

cache_key = id(node)

cached = self._scope_analysis_cache.get(cache_key)

if cached is not None:

return cached

collector = ScopeLocalNameCollector()

args_node = getattr(node, "args", None)

if args_node is not None:

collector.visit(args_node)

for statement in getattr(node, "body", []):

collector.visit(statement)

cached = {

"reserved_names": frozenset(collector.reserved_names),

"bindings": frozenset(collector.bindings - collector.external_bindings),

"external_bindings": frozenset(collector.external_bindings),

"args": frozenset(collector.args),

}

self._scope_analysis_cache[cache_key] = cached

return cached

def _scope_bindings(self, node):

analysis = self._scope_analysis(node)

return {

"bindings": set(analysis["bindings"]),

"globals": set(analysis["external_bindings"]),

"args": set(analysis["args"]),

}

def _is_preserved_public_global_reference(self, name):

if name not in self.public_global_names:

return False

for scope in reversed(self.scope_stack):

if name in scope["globals"]:

continue

if name in scope["bindings"]:

return False

return True

def _is_preserved_function_parameter_reference(self, node):

if self.rename_arguments:

return False

if self._is_in_function_signature(node):

return False

for scope in reversed(self.scope_stack):

if node.id in scope["globals"]:

continue

if node.id in scope["bindings"]:

return node.id in scope["args"]

return False

def _is_active_parameter_name(self, name):

for scope in reversed(self.scope_stack):

if name in scope["globals"]:

continue

if name in scope["bindings"]:

if name not in scope["args"]:

return False

if not self.rename_arguments:

return True

return name not in scope.get("renamed_args", set())

return False

def _local_scope_state(self, node):

analysis = self._scope_analysis(node)

reserved_names = set(analysis["reserved_names"])

local_bindings = analysis["bindings"]

for name in analysis["reserved_names"] - local_bindings:

visible_name = self._lookup_visible_identifier(name)

if visible_name is not None:

reserved_names.add(visible_name)

used_names = set(reserved_names)

return {

"mapping": {},

"used_names": used_names,

"generator": variable_name_generator(used_names),

}

def _receiver_class_name(self, node):

if isinstance(node, ast.Name):

return self._lookup_instance_type(node.id)

if isinstance(node, ast.Call):

func = node.func

if isinstance(func, ast.Name):

if func.id in self.class_member_mappings or func.id in self.class_method_argument_infos:

return func.id

class_context = self._current_class_context()

if class_context is not None and func.id in {

class_context["old_name"],

class_context["new_name"],

}:

return func.id

return None

def _record_instance_assignment(self, target, value):

class_name = self._receiver_class_name(value)

if isinstance(target, ast.Name):

self._set_instance_type(target.id, class_name)

return

if isinstance(target, (ast.Tuple, ast.List)):

for element in target.elts:

self._record_instance_assignment(element, value)

def _visit_ImportOrImportFrom(self, node):

"""Shorten imported library names.

>>> shortener = VariableShortener(variable_name_generator(), modules=('donotaliasme',))

>>> apply = lambda src: ast.unparse(shortener.visit(ast.parse(src)))

>>> apply('import demiurgic')

'import demiurgic as a'

>>> apply('from demiurgic import palpitation')

'from demiurgic import palpitation as b'

>>> apply('from demiurgic import a') # single-letter import should be left alone

'from demiurgic import a'

>>> print(apply('import demiurgic;demiurgic.palpitation()')) # TODO: bug - variable should remember object its bound to

import demiurgic as a

a.palpitation()

>>> print(apply('import demiurgic as dei;dei.palpitation()'))

import demiurgic as c

c.palpitation()

>>> print(apply('import demiurgic;import donotaliasme;from donotaliasme import dolor;'))

import demiurgic as a

import donotaliasme

from donotaliasme import dolor

"""

if self.keep_global_variables and self._is_node_global(node):

return self.generic_visit(node)

if isinstance(node, ast.Import) or node.module not in self.modules:

for alias in node.names:

if isinstance(node, ast.ImportFrom) or alias.name not in self.modules:

old = alias.asname or alias.name

if len(old) > 1:

alias.asname = self._rename_identifier(old)

return self.generic_visit(node)

visit_Import = _visit_ImportOrImportFrom

visit_ImportFrom = _visit_ImportOrImportFrom

def visit_ClassDef(self, node):

"""Shorten class names.

>>> shortener = VariableShortener(variable_name_generator())

>>> apply = lambda src: ast.unparse(shortener.visit(ast.parse(src)))

>>> apply('class Demiurgic: pass\\nholy = Demiurgic()')

'class a:\\n pass\\nb = a()'

>>> shortener = VariableShortener(variable_name_generator(), keep_global_variables=True)

>>> def apply(src):

... tree = ast.parse(src)

... shortener.visit(tree)

... append_public_aliases(tree, shortener.nodes_to_append)

... return ast.unparse(tree)

...

>>> apply('class Demiurgic: pass\\nholy = Demiurgic()')

'class Demiurgic:\\n pass\\nholy = Demiurgic()'

"""

old_name = node.name

parent_class_context = self._current_class_context()

rename_public_class = False

if self.keep_global_variables and self._is_node_global(node):

if (

len(node.name) > 1

and node.name not in self.mapping_values

and self._should_rename_public_class(node, old_name)

):

node.name = self._rename_identifier(old_name)

rename_public_class = old_name != node.name

class_context = {

"old_name": old_name,

"new_name": node.name,

"aliases": [],

"member_mapping": {},

"argument_infos": {},

"receiver_names": {"self", "cls"},

}

self.class_context_stack.append(class_context)

self.scope_stack.append(self._scope_bindings(node))

self._push_instance_scope()

try:

node = self.generic_visit(node)

finally:

self._pop_instance_scope()

self.scope_stack.pop()

self.class_context_stack.pop()

if class_context["member_mapping"]:

self.class_member_mappings[old_name] = dict(class_context["member_mapping"])

self.class_member_mappings[node.name] = dict(class_context["member_mapping"])

if class_context["argument_infos"]:

self.class_method_argument_infos[old_name] = dict(class_context["argument_infos"])

self.class_method_argument_infos[node.name] = dict(class_context["argument_infos"])

constructor_info = class_context["argument_infos"].get("__init__")

if constructor_info is not None:

copied = {

"rename_map": dict(constructor_info["rename_map"]),

"positional_params": list(constructor_info["positional_params"]),

}

self.callable_argument_infos[old_name] = copied

self.callable_argument_infos[node.name] = {

"rename_map": dict(copied["rename_map"]),

"positional_params": list(copied["positional_params"]),

}

if class_context["aliases"]:

node.body.extend(class_context["aliases"])

if rename_public_class:

return [

node,

self._generated_assignment(f"{old_name} = {node.name}"),

self._generated_assignment(f"{node.name}.__name__ = {old_name!r}"),

self._generated_assignment(f"{node.name}.__qualname__ = {old_name!r}"),

]

return node

if self.local_rename_scopes and not self._is_node_global(node):

node.name = self._rename_local_identifier(node.name)

elif node.name not in self.mapping_values:

node.name = self._rename_identifier(node.name)

if parent_class_context is not None and old_name != node.name:

parent_class_context["member_mapping"][old_name] = node.name

if self.keep_global_variables:

parent_class_context["aliases"].append(

self._generated_assignment(f"{old_name} = {node.name}")

)

class_context = {

"old_name": old_name,

"new_name": node.name,

"aliases": [],

"member_mapping": {},

"argument_infos": {},

"receiver_names": {"self", "cls"},

}

self.class_context_stack.append(class_context)

self.scope_stack.append(self._scope_bindings(node))

self._push_instance_scope()

try:

node = self.generic_visit(node)

finally:

self._pop_instance_scope()

self.scope_stack.pop()

self.class_context_stack.pop()

if class_context["member_mapping"]:

self.class_member_mappings[old_name] = dict(class_context["member_mapping"])

self.class_member_mappings[node.name] = dict(class_context["member_mapping"])

if class_context["argument_infos"]:

self.class_method_argument_infos[old_name] = dict(class_context["argument_infos"])

self.class_method_argument_infos[node.name] = dict(class_context["argument_infos"])

constructor_info = class_context["argument_infos"].get("__init__")

if constructor_info is not None:

copied = {

"rename_map": dict(constructor_info["rename_map"]),

"positional_params": list(constructor_info["positional_params"]),

}

self.callable_argument_infos[old_name] = copied

self.callable_argument_infos[node.name] = {

"rename_map": dict(copied["rename_map"]),

"positional_params": list(copied["positional_params"]),

}

return node

def visit_FunctionDef(self, node):

"""Shorten function names.

>>> shortener = VariableShortener(variable_name_generator())

>>> apply = lambda src: ast.unparse(shortener.visit(ast.parse(src)))

>>> apply('def demiurgic(palpitation): return palpitation\\nholy = demiurgic()')

'def a(palpitation):\\n return palpitation\\nb = a()'

>>> shortener = VariableShortener(variable_name_generator(), keep_global_variables=True)

>>> def apply(src):

... tree = ast.parse(src)

... shortener.visit(tree)

... append_public_aliases(tree, shortener.nodes_to_append)

... return ast.unparse(tree)

...

>>> apply('def demiurgic(palpitation): return palpitation\\nholy = demiurgic()')

'def a(palpitation):\\n return palpitation\\nholy = a()\\ndemiurgic = a'

"""

old_name = node.name

if self._preserve_function_name(node.name):

self.scope_stack.append(self._scope_bindings(node))

self.local_rename_scopes.append(self._local_scope_state(node))

self._push_instance_scope()

try:

argument_info = self._rename_function_arguments(node)

self.scope_stack[-1]["renamed_args"] = set(argument_info["rename_map"])

class_context = self._current_class_context()

if class_context is not None:

class_context["receiver_names"].update(argument_info["receiver_names"])

if argument_info["rename_map"] or argument_info["positional_params"]:

copied = {

"rename_map": dict(argument_info["rename_map"]),

"positional_params": list(argument_info["positional_params"]),

}

class_context["argument_infos"][old_name] = copied

return self.generic_visit(node)

finally:

self._pop_instance_scope()

self.local_rename_scopes.pop()

self.scope_stack.pop()

if self._is_method_definition(node):

class_context = self._current_class_context()

class_name = class_context["old_name"] if class_context is not None else ""

if (

len(old_name) > 1

and (

not self.keep_global_variables

or class_context is None

or self._should_rename_public_member(node.parent, class_name, old_name)

)

):

node.name = self._rename_identifier(old_name)

if class_context is not None and old_name != node.name:

class_context["member_mapping"][old_name] = node.name

if self.keep_global_variables:

class_context["aliases"].append(

self._generated_assignment(f"{old_name} = {node.name}")

)

self.scope_stack.append(self._scope_bindings(node))

self.local_rename_scopes.append(self._local_scope_state(node))

self._push_instance_scope()

try:

argument_info = self._rename_function_arguments(node)

self.scope_stack[-1]["renamed_args"] = set(argument_info["rename_map"])