# -------------------------------------------------------------------------- # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. # -------------------------------------------------------------------------- # pylint: disable=W0613 from __future__ import annotations import io from abc import ABC, abstractmethod from enum import Enum from textwrap import TextWrapper, dedent from typing import ( Any, Callable, Generic, Iterable, Optional, Set, TextIO, Tuple, Type, TypeVar, Union, ) T = TypeVar("T") TNode = TypeVar("TNode", bound="Node") TExpr = TypeVar("TExpr", bound="Expr") NoneType = type(None) def _assert_instance(instance, expected_type: Union[Type, Tuple[Type, ...]]): if not isinstance(instance, expected_type): raise TypeError(f"expected: {expected_type!r}; actual: {instance!r}") __end_of_sequence = StopIteration() def first_or_none(seq: Iterable[T]) -> Optional[T]: return next(iter(seq), None) def first(seq: Iterable[T]) -> T: return next(iter(seq)) def single_or_none(seq: Iterable[T]) -> Optional[T]: i = iter(seq) value = next(i, __end_of_sequence) if value is __end_of_sequence: return None if next(i, __end_of_sequence) is not __end_of_sequence: raise StopIteration("sequence contains more than one element") return value class Role: def __init__(self, name: str): _assert_instance(name, str) self.name = name def __str__(self): return self.name class NodePredicate: always: NodePredicate def __init__( self, role: Optional[Role] = None, type_: Optional[Type[TNode]] = None, func: Optional[Callable[[Node], bool]] = None, ): _assert_instance(role, (Role, NoneType)) _assert_instance(type_, (type, NoneType)) self.role = role self.type = type_ self.func = func def matches(self, node: Node): _assert_instance(node, Node) matches = True if self.role: matches &= node.role is self.role if self.type: matches &= isinstance(node, self.type) if self.func and matches: matches &= self.func(node) return matches NodePredicate.always = NodePredicate() class Node(ABC): # pylint: disable=W0212 def __init__(self): self._role: Optional[Role] = None self._parent: Optional[Node] = None self._prev_sibling: Optional[Node] = None self._next_sibling: Optional[Node] = None self._first_child: Optional[Node] = None self._last_child: Optional[Node] = None self.leading_trivia: Optional[str] = None self.trailing_trivia: Optional[str] = None @property def qual_name(self) -> str: names = [] for ancestor in self.get_ancestors(and_self=True): names.insert(0, ancestor.name if hasattr(ancestor, "name") else "") return ".".join(names) @property def parent_module(self) -> Optional[Module]: return first_or_none(self.get_ancestors_of_type(Module)) @property def parent(self): return self._parent @property def role(self): return self._role @property def prev_sibling(self): return self._prev_sibling @property def next_sibling(self): return self._next_sibling @property def first_child(self): return self._first_child @property def last_child(self): return self._last_child @property def has_children(self): return self._first_child is not None @property def children(self) -> Iterable[Node]: current_node = self.first_child while current_node is not None: # save next then yield to allow removing/replacing nodes while iterating next_node = current_node.next_sibling yield current_node current_node = next_node def get_children(self, predicate: NodePredicate) -> Iterable[Node]: _assert_instance(predicate, NodePredicate) yield from filter(predicate.matches, self.children) def get_children_in_role(self, role: Role): _assert_instance(role, Role) return self.get_children(NodePredicate(role=role)) def get_children_of_type(self, type_: Type[TNode]) -> Iterable[TNode]: _assert_instance(type_, type) return self.get_children(NodePredicate(type_=type_)) def get_ancestors( self, predicate: Optional[NodePredicate] = None, and_self=False ) -> Iterable[Node]: current_node = self if and_self else self.parent while current_node: # save next then yield to allow removing/replacing nodes while iterating next_node = current_node.parent if predicate is None or predicate.matches(current_node): yield current_node current_node = next_node def get_ancestors_in_role(self, role: Role, and_self=False): _assert_instance(role, Role) return self.get_ancestors(NodePredicate(role=role), and_self=and_self) def get_ancestors_of_type(self, type_: Type[TNode], and_self=False) -> Iterable[TNode]: _assert_instance(type_, type) return self.get_ancestors(NodePredicate(type_=type_), and_self=and_self) def _set_parent(self, child: Node): if child._parent is not None: raise ValueError(f"node is already has a parent: {child.parent!r}") child._parent = self def _get_single_child(self, role: Role) -> Optional[Node]: return first_or_none(self.get_children_in_role(role)) def _set_single_child(self, node: Node, role: Role): current_node = self._get_single_child(role) if current_node: current_node.replace(node) else: self.append_child(node, role) def append_children(self, children: Optional[Union[Node, Iterable[Node]]], role: Role): _assert_instance(role, Role) if children is None: return if isinstance(children, Node): self.append_child(children, role) else: for child in children: self.append_child(child, role) def append_child(self, child: Node, role: Role): _assert_instance(role, Role) if child is None: return _assert_instance(child, Node) self._set_parent(child) child._role = role if self._first_child is None: self._last_child = child self._first_child = child else: self._last_child._next_sibling = child child._prev_sibling = self._last_child self._last_child = child def insert_child_before(self, next_sibling: Optional[Node], child: Node, role: Role): _assert_instance(next_sibling, (Node, type(None))) _assert_instance(child, Node) _assert_instance(role, Role) if next_sibling is None: self.append_child(child, role) return self._set_parent(child) child._role = role child._next_sibling = next_sibling child._prev_sibling = next_sibling._prev_sibling if next_sibling._prev_sibling is None: self._first_child = child else: next_sibling._prev_sibling._next_sibling = child next_sibling._prev_sibling = child def prepend_child(self, child: Node, role: Role): _assert_instance(child, Node) _assert_instance(role, Role) self.insert_child_before(self.first_child, child, role) def remove(self): if self._prev_sibling is not None: self._prev_sibling._next_sibling = self._next_sibling else: self._parent._first_child = self._next_sibling if self._next_sibling is not None: self._next_sibling._prev_sibling = self._prev_sibling else: self._parent._last_child = self._prev_sibling self._parent = None self._role = None self._prev_sibling = None self._next_sibling = None return self def replace(self, new_node: Optional[Node]): if new_node is None: self.remove() return if new_node is self: return if self.parent is None: raise ValueError("cannot replace root node") _assert_instance(new_node, Node) if new_node.parent is not None: if self in new_node.ancestors: new_node.remove() else: raise ValueError(f"node is used in another tree: {new_node!r}") new_node._parent = self._parent new_node._role = self._role new_node._prev_sibling = self._prev_sibling new_node._next_sibling = self._next_sibling if self._prev_sibling is None: self._parent._first_child = new_node else: self._prev_sibling._next_sibling = new_node if self._next_sibling is None: self._parent._last_child = new_node else: self._parent._prev_sibling = new_node self._parent = None self._role = None self._prev_sibling = None self._next_sibling = None @abstractmethod def accept(self, visitor: Visitor): pass def _dispatch_visit(self, dispatch: Callable[[TNode, VisitKind], bool]): visitor = dispatch.__self__ visitor.enter(self) if dispatch(self) is True: for child in self.children: child.accept(dispatch.__self__) visitor.leave(self) dispatch(self) else: visitor.leave(self) visitor.finish(self) def __str__(self): buffer = io.StringIO() self.accept(PythonWriter(buffer)) return buffer.getvalue() class Expr(Node, ABC): def accept(self, visitor: Visitor): self._dispatch_visit(visitor.visit_expr) class ThunkExpr(Expr): def __init__(self, code: str): super().__init__() self.code = code def accept(self, visitor: Visitor): self._dispatch_visit(visitor.visit_thunk_expr) class Name(Expr): def __init__(self, identifier: str): super().__init__() self.identifier = identifier def accept(self, visitor: Visitor): self._dispatch_visit(visitor.visit_name) class Constant(Expr): def __init__(self, value: Any): super().__init__() self.value = value def accept(self, visitor: Visitor): self._dispatch_visit(visitor.visit_constant) class ExprList(Expr, ABC, Generic[TExpr]): class Roles: Elements = Role("ExprList.Elements") def __init__(self, *elements: TExpr): super().__init__() self.append_children(elements, ExprList.Roles.Elements) @property def elements(self) -> Iterable[TExpr]: return self.get_children_in_role(ExprList.Roles.Elements) def append_element(self, element: TExpr): _assert_instance(element, Expr) self.append_child(element, ExprList.Roles.Elements) class BinOp(Expr): class Roles: Left = Role("BinOp.Left") Right = Role("BinOp.Right") def __init__(self, left: Expr, op: str, right: Expr): super().__init__() self.append_child(left, BinOp.Roles.Left) self.op = op self.append_child(right, BinOp.Roles.Right) @property def left(self): return first(self.get_children_in_role(BinOp.Roles.Left)) @property def right(self): return first(self.get_children_in_role(BinOp.Roles.Right)) def accept(self, visitor: Visitor): self._dispatch_visit(visitor.visit_binop) class Subscript(Expr): class Roles: Value = Role("Subscript.Value") Slice = Role("Subscript.Slice") def __init__(self, value: Expr, slice: Expr): super().__init__() self.append_child(value, Subscript.Roles.Value) self.append_child(slice, Subscript.Roles.Slice) @property def value(self): return first(self.get_children_in_role(Subscript.Roles.Value)) @property def slice(self): return first(self.get_children_in_role(Subscript.Roles.Slice)) def accept(self, visitor: Visitor): self._dispatch_visit(visitor.visit_subscript) class Starred(Expr): class Roles: Expr = Role("Starred.Expr") def __init__(self, expr: Expr): super().__init__() self.append_child(expr, Starred.Roles.Expr) @property def expr(self): return first(self.get_children_in_role(Starred.Roles.Expr)) def accept(self, visitor: Visitor): self._dispatch_visit(visitor.visit_starred) class Call(Expr): class Roles: Func = Role("Call.Func") Args = Role("Call.Args") def __init__(self, func: Expr, *args: Expr): super().__init__() _assert_instance(func, Expr) self.append_child(func, Call.Roles.Func) self.append_children(args, Call.Roles.Args) @property def func(self) -> Expr: return first(self.get_children_in_role(Call.Roles.Func)) @property def args(self) -> Iterable[Expr]: return self.get_children_in_role(Call.Roles.Args) def accept(self, visitor: Visitor): self._dispatch_visit(visitor.visit_call) class Lambda(Expr): class Roles: Args = Role("Lambda.Args") Body = Role("Lambda.Body") def __init__(self, body: Expr, *args: Arg): super().__init__() _assert_instance(body, Expr) self.append_child(body, Lambda.Roles.Body) self.append_children(args, Lambda.Roles.Args) @property def body(self) -> Expr: return first(self.get_children_in_role(Lambda.Roles.Body)) @property def args(self) -> Iterable[Expr]: return self.get_children_in_role(Lambda.Roles.Args) def accept(self, visitor: Visitor): self._dispatch_visit(visitor.visit_lambda) class TupleExpr(ExprList): def accept(self, visitor: Visitor): self._dispatch_visit(visitor.visit_tuple_expr) class ListExpr(ExprList): def accept(self, visitor: Visitor): self._dispatch_visit(visitor.visit_list_expr) class SetExpr(ExprList): def accept(self, visitor: Visitor): self._dispatch_visit(visitor.visit_set_expr) class DictElem(Expr): class Roles: Key = Role("DictElem.Key") Value = Role("DictElem.Value") def __init__(self, key: Expr, value: Expr): super().__init__() _assert_instance(key, Expr) _assert_instance(value, Expr) self.append_child(key, DictElem.Roles.Key) self.append_child(value, DictElem.Roles.Value) @property def key(self) -> Expr: return first(self.get_children_in_role(DictElem.Roles.Key)) @property def value(self) -> Expr: return first(self.get_children_in_role(DictElem.Roles.Value)) def accept(self, visitor: Visitor): self._dispatch_visit(visitor.visit_dict_elem) class DictExpr(ExprList[DictElem]): def accept(self, visitor: Visitor): self._dispatch_visit(visitor.visit_dict_expr) class TypeRef(Expr): class Roles: TypeArgs = Role("TypeRef.TypeArgs") def __init__( self, module: Optional[str], name: str, *typeargs: TypeRef, default_value: Optional[Constant] = None, ): super().__init__() self.module = module self.name = name self.default_value = default_value or Constant(None) self.imported_by: Optional[ImportBase] = None self.append_children(typeargs, TypeRef.Roles.TypeArgs) @property def typeargs(self) -> Iterable[TypeRef]: return self.get_children_in_role(TypeRef.Roles.TypeArgs) def append_typearg(self, typearg: TypeRef): _assert_instance(typearg, TypeRef) self.append_child(typearg, TypeRef.Roles.TypeArgs) def accept(self, visitor: Visitor): self._dispatch_visit(visitor.visit_typeref) @staticmethod def make_composite_if_multiple( composite_type: type[TypeRef], *typeargs: TypeRef ) -> TypeRef: if len(typeargs) == 0: return NoneTypeRef elif len(typeargs) == 1: return typeargs[0] else: return composite_type(*typeargs) class BuiltinTypeRef(TypeRef): def __init__(self, name: str, *typeargs: TypeRef, **kwargs): super().__init__(None, name, *typeargs, **kwargs) class NoneTypeRef(BuiltinTypeRef): def __init__(self): super().__init__("None") class BoolTypeRef(BuiltinTypeRef): def __init__(self): super().__init__("bool", default_value=Constant(bool())) # noqa: UP018 class IntTypeRef(BuiltinTypeRef): def __init__(self): super().__init__("int", default_value=Constant(int())) # noqa: UP018 class FloatTypeRef(BuiltinTypeRef): def __init__(self): super().__init__("float", default_value=Constant(float())) # noqa: UP018 class ComplexTypeRef(BuiltinTypeRef): def __init__(self): super().__init__("complex", default_value=Constant(complex())) class StrTypeRef(BuiltinTypeRef): def __init__(self): super().__init__("str") class BytesTypeRef(BuiltinTypeRef): def __init__(self): super().__init__("bytes") class EllipsisTypeRef(BuiltinTypeRef): def __init__(self): super().__init__("...") class TypingRefs(ABC): @abstractmethod def __init__(self): pass class Any(TypeRef): def __init__(self): super().__init__("typing", "Any") class Union(TypeRef): def __init__(self, *typeargs: TypeRef): super().__init__("typing", "Union", *typeargs) class Optional(TypeRef): def __init__(self, *typeargs: TypeRef): super().__init__("typing", "Optional", *typeargs) class Sequence(TypeRef): def __init__(self, *typeargs: TypeRef): super().__init__("typing", "Sequence", *typeargs) class Tuple(TypeRef): def __init__(self, *typeargs: TypeRef): super().__init__("typing", "Tuple", *typeargs) class Mapping(TypeRef): def __init__(self, *typeargs: TypeRef): super().__init__("typing", "Mapping", *typeargs) class List(TypeRef): def __init__(self, *typeargs: TypeRef): super().__init__("typing", "List", *typeargs) class Annotation(TypeRef): def __init__(self, *typeargs: TypeRef): super().__init__("typing", "Annotation", *typeargs) class Callable(TypeRef): def __init__(self, *typeargs: TypeRef): super().__init__("typing", "Callable", *typeargs) class Arg(Node): class Roles: Type = Role("Arg.Type") DefaultValue = Role("Arg.DefaultValue") def __init__( self, name: str, type: Optional[TypeRef] = None, default_value: Optional[Expr] = None, is_vararg: bool = False, is_kwarg: bool = False, doc: Optional[str] = None, ): super().__init__() self.name = name self.is_vararg = is_vararg self.is_kwarg = is_kwarg self.doc = doc self.append_child(type, Arg.Roles.Type) self.append_child(default_value, Arg.Roles.DefaultValue) @property def type(self) -> Optional[TypeRef]: return first_or_none(self.get_children_in_role(Arg.Roles.Type)) @property def default_value(self) -> Optional[Expr]: return first_or_none(self.get_children_in_role(Arg.Roles.DefaultValue)) @default_value.setter def default_value(self, value: Optional[Expr]): self._set_single_child(value, Arg.Roles.DefaultValue) @property def has_default_value(self) -> bool: return self.default_value is not None def accept(self, visitor: Visitor): self._dispatch_visit(visitor.visit_arg) class Stmt(Node, ABC): pass class BlockStmt(Stmt, ABC): pass class Pass(Stmt): def accept(self, visitor: Visitor): self._dispatch_visit(visitor.visit_pass) class ThunkStmt(Stmt): class Roles: Thunk = Role("ThunkStmt.Thunk") def __init__(self, *thunks: Union[str, Stmt]): super().__init__() self.thunk: Optional[str] = None if len(thunks) == 1 and isinstance(thunks[0], str): self.thunk = thunks[0] else: for thunk in thunks: if isinstance(thunk, str): self.append_child(ThunkStmt(thunk), ThunkStmt.Roles.Thunk) else: self.append_child(thunk, ThunkStmt.Roles.Thunk) def accept(self, visitor: Visitor): self._dispatch_visit(visitor.visit_thunk_stmt) class FunctionDef(BlockStmt): class Roles: Args = Role("FunctionDef.Args") ReturnType = Role("FunctionDef.ReturnType") Body = Role("FunctionDef.Body") def __init__( self, name: str, *args: Arg, return_type: Optional[TypeRef] = None, body: Union[Stmt, Iterable[Stmt]] = (), doc: Optional[str] = None, ): super().__init__() self.name = name self.doc = doc self.append_children(args, FunctionDef.Roles.Args) self.append_children(return_type, FunctionDef.Roles.ReturnType) self.append_children(body, FunctionDef.Roles.Body) @property def args(self) -> Iterable[Arg]: return self.get_children_in_role(FunctionDef.Roles.Args) def append_arg(self, base: TypeRef): _assert_instance(base, TypeRef) self.append_child(base, FunctionDef.Roles.Args) @property def return_type(self) -> Optional[TypeRef]: return self._get_single_child(FunctionDef.Roles.ReturnType) @return_type.setter def return_type(self, return_type: Optional[TypeRef]): self._set_single_child(return_type, FunctionDef.Roles.ReturnType) @property def body(self) -> Iterable[Stmt]: return self.get_children_in_role(FunctionDef.Roles.Body) def append_body(self, stmt: Stmt): _assert_instance(stmt, Stmt) self.append_child(stmt, FunctionDef.Roles.Body) def accept(self, visitor: Visitor): self._dispatch_visit(visitor.visit_functiondef) class ClassDef(BlockStmt): class Roles: Bases = Role("ClassDef.Bases") Body = Role("ClassDef.Body") def __init__(self, name: str, *body: Stmt, bases: Union[TypeRef, Iterable[TypeRef]] = ()): super().__init__() self.name = name self.append_children(bases, ClassDef.Roles.Bases) self.append_children(body, ClassDef.Roles.Body) @property def bases(self) -> Iterable[TypeRef]: return self.get_children_in_role(ClassDef.Roles.Bases) def append_base(self, base: TypeRef): _assert_instance(base, TypeRef) self.append_child(base, ClassDef.Roles.Bases) @property def body(self) -> Iterable[Stmt]: return self.get_children_in_role(ClassDef.Roles.Body) def make_typeref(self) -> TypeRef: return TypeRef(self.parent.qual_name if self.parent else None, self.name) def append_body(self, stmt: Stmt): _assert_instance(stmt, Stmt) self.append_child(stmt, ClassDef.Roles.Body) def accept(self, visitor: Visitor): self._dispatch_visit(visitor.visit_classdef) class Return(Stmt): class Roles: Expr = Role("Return.Expr") def __init__(self, expr: Expr): super().__init__() self.append_child(expr, Return.Roles.Expr) @property def expr(self): return self._get_single_child(Return.Roles.Expr) def accept(self, visitor: Visitor): self._dispatch_visit(visitor.visit_return) class Assign(Stmt): class Roles: Target = Role("Assign.Target") Value = Role("Assign.Value") Type = Role("Assign.Type") def __init__(self, target: Expr, value: Expr, type: Optional[TypeRef] = None): super().__init__() self.target = target self.value = value self.type = type @property def target(self) -> Optional[Expr]: return self._get_single_child(Assign.Roles.Target) @target.setter def target(self, expr: Optional[Expr]): self._set_single_child(expr, Assign.Roles.Target) @property def value(self) -> Optional[Expr]: return self._get_single_child(Assign.Roles.Value) @value.setter def value(self, expr: Optional[Expr]): self._set_single_child(expr, Assign.Roles.Value) @property def type(self) -> Optional[TypeRef]: return self._get_single_child(Assign.Roles.Type) @type.setter def type(self, expr: Optional[TypeRef]): self._set_single_child(expr, Assign.Roles.Type) def accept(self, visitor: Visitor): self._dispatch_visit(visitor.visit_assign) class If(BlockStmt): class Roles: Condition = Role("If.Condition") TrueBody = Role("If.TrueBody") FalseBody = Role("If.FalseBody") def __init__( self, condition: Expr, true_body: Iterable[Stmt], false_body: Optional[Iterable[Stmt]] = None, ): super().__init__() self.condition = condition self.append_children(true_body, If.Roles.TrueBody) self.append_children(false_body, If.Roles.FalseBody) @property def condition(self) -> Optional[Expr]: return self._get_single_child(If.Roles.Condition) @condition.setter def condition(self, expr: Optional[Expr]): self._set_single_child(expr, If.Roles.Condition) @property def true_body(self) -> Iterable[Stmt]: return self.get_children_in_role(If.Roles.TrueBody) @property def false_body(self) -> Iterable[Stmt]: return self.get_children_in_role(If.Roles.FalseBody) def accept(self, visitor: Visitor): self._dispatch_visit(visitor.visit_if) class Raise(Node): class Roles: Expr = Role("Raise.Expr") def __init__(self, expr: Expr): super().__init__() self.append_child(expr, Raise.Roles.Expr) @property def expr(self): return first(self.get_children_in_role(Raise.Roles.Expr)) def accept(self, visitor: Visitor): self._dispatch_visit(visitor.visit_raise) class Alias(Node): def __init__(self, name: str, alias: Optional[str] = None): super().__init__() self.name = name self.alias = alias def accept(self, visitor: Visitor): self._dispatch_visit(visitor.visit_alias) class ImportBase(Stmt, ABC): class Roles: Names = Role("ImportBase.Names") def __init__(self, *names: Alias): super().__init__() self.append_children(names, ImportBase.Roles.Names) @property def names(self) -> Iterable[Alias]: return self.get_children_in_role(ImportBase.Roles.Names) class Import(ImportBase): def accept(self, visitor: Visitor): self._dispatch_visit(visitor.visit_import) class ImportFrom(ImportBase): def __init__(self, module: str, *names: Alias, level: Optional[int] = None): super().__init__(*names) self.module = module self.level = level def accept(self, visitor: Visitor): self._dispatch_visit(visitor.visit_importfrom) class Module(Node): class Roles: Body = Role("Module.Body") def __init__(self, *body: Stmt, name: Optional[str] = None): super().__init__() self.name = name self.append_children(body, Module.Roles.Body) @property def body(self) -> Iterable[Stmt]: return self.get_children_in_role(Module.Roles.Body) def append_body(self, *stmts: Node): self.append_children(stmts, Module.Roles.Body) def accept(self, visitor: Visitor): self._dispatch_visit(visitor.visit_module) class VisitKind(Enum): NONE = 0 ENTER = 1 LEAVE = 2 class Visitor: def __init__(self): self.visit_kind = VisitKind.NONE self.node_stack = [] def enter(self, node: Node): self.visit_kind = VisitKind.ENTER self.node_stack.append(node) def leave(self, node: Node): self.visit_kind = VisitKind.LEAVE def finish(self, node: Node): self.visit_kind = VisitKind.NONE self.node_stack.pop() def visit_node(self, node: Node) -> Optional[bool]: return True def visit_expr(self, expr: Expr) -> Optional[bool]: return self.visit_node(expr) def visit_name(self, name: Name) -> Optional[bool]: return self.visit_expr(name) def visit_constant(self, constant: Constant) -> Optional[bool]: return self.visit_expr(constant) def visit_binop(self, binop: BinOp) -> Optional[bool]: return self.visit_expr(binop) def visit_subscript(self, subscript: Subscript) -> Optional[bool]: return self.visit_expr(subscript) def visit_starred(self, starred: Starred) -> Optional[bool]: return self.visit_expr(starred) def visit_call(self, call: Call) -> Optional[bool]: return self.visit_expr(call) def visit_lambda(self, lambda_: Lambda) -> Optional[bool]: return self.visit_expr(lambda_) def visit_expr_list(self, expr_list: ExprList) -> Optional[bool]: return self.visit_expr(expr_list) def visit_thunk_expr(self, thunk: ThunkExpr) -> Optional[bool]: return self.visit_expr(thunk) def visit_tuple_expr(self, tuple: TupleExpr) -> Optional[bool]: return self.visit_expr_list(tuple) def visit_list_expr(self, list: ListExpr) -> Optional[bool]: return self.visit_expr_list(list) def visit_set_expr(self, set: SetExpr) -> Optional[bool]: return self.visit_expr_list(set) def visit_dict_elem(self, elem: DictElem) -> Optional[bool]: return self.visit_expr(elem) def visit_dict_expr(self, dict: DictExpr) -> Optional[bool]: return self.visit_expr_list(dict) def visit_typeref(self, typeref: TypeRef) -> Optional[bool]: return self.visit_expr(typeref) def visit_arg(self, arg: Arg) -> Optional[bool]: return self.visit_node(arg) def visit_stmt(self, stmt: Stmt) -> Optional[bool]: return self.visit_node(stmt) def visit_blockstmt(self, block: BlockStmt) -> Optional[bool]: return self.visit_stmt(block) def visit_pass(self, pass_: Pass) -> Optional[bool]: return self.visit_stmt(pass_) def visit_thunk_stmt(self, thunk: ThunkStmt) -> Optional[bool]: return self.visit_stmt(thunk) def visit_functiondef(self, functiondef: FunctionDef) -> Optional[bool]: return self.visit_stmt(functiondef) def visit_classdef(self, classdef: ClassDef) -> Optional[bool]: return self.visit_stmt(classdef) def visit_return(self, return_: Return) -> Optional[bool]: return self.visit_stmt(return_) def visit_assign(self, assign: Assign) -> Optional[bool]: return self.visit_stmt(assign) def visit_if(self, if_: If) -> Optional[bool]: return self.visit_stmt(if_) def visit_raise(self, raise_: Raise) -> Optional[bool]: return self.visit_stmt(raise_) def visit_alias(self, alias: Alias) -> Optional[bool]: return self.visit_node(alias) def visit_importbase(self, import_: ImportBase) -> Optional[bool]: return self.visit_stmt(import_) def visit_import(self, import_: Import) -> Optional[bool]: return self.visit_importbase(import_) def visit_importfrom(self, importfrom: ImportFrom) -> Optional[bool]: return self.visit_importbase(importfrom) def visit_module(self, module: Module) -> Optional[bool]: return self.visit_node(module) class FixupVisitor(Visitor, ABC): pass class PopulateEmptyMemberBodies(FixupVisitor): def visit_classdef(self, classdef: ClassDef) -> Optional[bool]: if self.visit_kind is VisitKind.ENTER and not any(classdef.body): classdef.append_child(Pass(), ClassDef.Roles.Body) return True def visit_functiondef(self, functiondef: FunctionDef) -> Optional[bool]: if self.visit_kind is VisitKind.ENTER and not any(functiondef.body): functiondef.append_child(Pass(), FunctionDef.Roles.Body) return True class NameCollector(Visitor): def __init__(self, predicate: NodePredicate): super().__init__() _assert_instance(predicate, NodePredicate) self._predicate = predicate self.names: Set[str] = set() def leave(self, node: Node) -> Optional[bool]: if self._predicate.matches(node) and hasattr(node, "name"): self.names.add(node.name) class ImportAdjuster(FixupVisitor): def __init__(self): super().__init__() self.naming_conflicts: Set[str] = set() def enter(self, node: Node): if len(self.node_stack) == 0: collector = NameCollector( NodePredicate(func=lambda n: isinstance(n, (ClassDef, FunctionDef))) ) node.accept(collector) self.naming_conflicts = collector.names super().enter(node) def leave(self, node: Node): super().leave(node) if len(self.node_stack) == 0: self.naming_conflicts = set() def visit_typeref(self, typeref: TypeRef) -> Optional[bool]: if self.visit_kind is not VisitKind.ENTER or not typeref.module: return True module = first_or_none(typeref.get_ancestors_of_type(Module)) if module is None: return True def adjust_typeref(import_alias: Optional[str]): typeref.module = None if import_alias: typeref.name = import_alias import_from: ImportFrom = None # Reuse an existing import if we have one; if so, # and the imported name is already specified, return # early as there's nothing to import. In that case, also # adjust the typeref if the import is aliased due to # conflict resolution below from a previous pass. for import_ in filter( lambda i: i.module == typeref.module, module.get_children_of_type(ImportFrom) ): import_from = import_ for imported_name in filter( lambda i: i.name in (typeref.name, typeref.name), import_.names ): adjust_typeref(imported_name.alias) return True # See if the type name conflicts with other names in the # module (class and function names). If so, adjust the # name to create an alias on the import. This rewrites # conflicts like: # from typing import Optional # def Optional(thing: Optional[str]): ... # To: # from typing import Optional as _Optional # def Optional(thing: _Optional[str]): ... conflict_alias = typeref.name while conflict_alias in self.naming_conflicts: conflict_alias = f"_{conflict_alias}" if conflict_alias == typeref.name: import_alias = Alias(typeref.name) else: import_alias = Alias(typeref.name, conflict_alias) # Expand or create the import if import_from is None: module.prepend_child(ImportFrom(typeref.module, import_alias), Module.Roles.Body) else: import_from.append_child(import_alias, ImportBase.Roles.Names) adjust_typeref(conflict_alias) return True class NodeWriterOptions: def __init__(self, indent=" ", newline="\n", insert_final_newline=True): self.indent = indent self.newline = newline self.insert_final_newline = insert_final_newline class NodeWriter(Visitor, ABC): def __init__(self, stream: TextIO, options: Optional[NodeWriterOptions] = None): super().__init__() self._stream = stream self._options = options or NodeWriterOptions() self._indent_level = 0 self._last_char = "" def enter(self, node: Node): super().enter(node) if node.leading_trivia: self.write(node.leading_trivia) def finish(self, node: Node): super().finish(node) if node.trailing_trivia: self.write(node.trailing_trivia) if self._options.insert_final_newline and len(self.node_stack) == 0: self.write("\n") def indent(self): self._indent_level += 1 def dedent(self): self._indent_level -= 1 def write_indent(self): self._stream.write(self._options.indent * self._indent_level) def _raw_write(self, str: str): if len(str) > 0: if self._options.newline != "\n": self._stream.write(str.replace("\n", self._options.newline)) else: self._stream.write(str) self._last_char = str[-1] def write(self, *texts: str, separator: str = "", allow_empty_text: bool = False): for i, text in enumerate(texts): if not allow_empty_text and len(text) == 0: continue if self._last_char == "\n": self.write_indent() if i > 0: self._raw_write(separator) if separator == "\n": self.write_indent() self._raw_write(text) def dispatch_write( self, separator: Union[str, Callable[[Node], str]], nodes: Iterable[Node], prefix: str = "", suffix: str = "", ): self.write(prefix) for i, node in enumerate(nodes): if i > 0: if callable(separator): self.write(separator(node)) else: self.write(separator) node.accept(self) self.write(suffix) class PythonWriter(NodeWriter): def visit_node(self, node: Node) -> Optional[bool]: raise NotImplementedError(f"no visitor for node {node}") def visit_module(self, module: Module): def sep(node: Node): node_is_block = isinstance(node, BlockStmt) prev_is_block = isinstance(node.prev_sibling, BlockStmt) if prev_is_block or (node_is_block and not prev_is_block): return "\n\n\n" else: return "\n" self.dispatch_write(sep, module.body) def visit_alias(self, alias: Alias): self.write(alias.name) if alias.alias: self.write(" as ") self.write(alias.alias) def visit_import(self, import_: Import): self.write("import ") self.dispatch_write(", ", import_.names) def visit_importfrom(self, importfrom: ImportFrom): self.write(f"from {importfrom.module} import ") self.dispatch_write(", ", importfrom.names) def visit_typeref(self, typeref: TypeRef): if typeref.module and len(typeref.module) > 0: self.write(typeref.module) self.write(".") self.write(typeref.name) if any(typeref.typeargs): self.write("[") self.dispatch_write(", ", typeref.typeargs) self.write("]") def visit_arg(self, arg: Arg): if arg.is_vararg: self.write("*") self.write(arg.name) if arg.type: self.write(": ") arg.type.accept(self) if arg.default_value: self.write(" = ") arg.default_value.accept(self) def visit_thunk_expr(self, thunk: ThunkExpr): self.write(thunk.code) def visit_name(self, name: Name): self.write(name.identifier) def visit_constant(self, constant: Constant): self.write( repr(constant.value) if isinstance(constant.value, str) else str(constant.value) ) def visit_binop(self, binop: BinOp): binop.left.accept(self) self.write(f" {binop.op} ") binop.right.accept(self) def visit_subscript(self, subscript: Subscript): subscript.value.accept(self) self.write("[") subscript.slice.accept(self) self.write("]") def visit_starred(self, starred: Starred): self.write("*") starred.expr.accept(self) def visit_call(self, call: Call): call.func.accept(self) self.dispatch_write(", ", call.args, prefix="(", suffix=")") def visit_lambda(self, lambda_: Lambda): self.write("lambda ") self.dispatch_write(", ", lambda_.args) self.write(": ") lambda_.body.accept(self) def visit_tuple_expr(self, tuple: TupleExpr): self.dispatch_write(", ", tuple.elements, prefix="(", suffix=",)") def visit_list_expr(self, list: ListExpr): self.dispatch_write(", ", list.elements, prefix="[", suffix="]") def visit_set_expr(self, set: ListExpr): self.dispatch_write(", ", set.elements, prefix="{", suffix="}") def visit_dict_elem(self, elem: DictElem): elem.key.accept(self) self.write(": ") elem.value.accept(self) def visit_dict_expr(self, dict: DictExpr): self.dispatch_write(", ", dict.elements, prefix="{", suffix="}") def visit_pass(self, pass_: Pass): self.write("pass") def visit_thunk_stmt(self, thunk: ThunkStmt) -> bool: if self.visit_kind == VisitKind.ENTER and thunk.thunk: lines = dedent(thunk.thunk).splitlines() self.write(*lines, separator="\n", allow_empty_text=True) if thunk.next_sibling: self.write("\n") return True def visit_assign(self, assign: Assign): assign.target.accept(self) if assign.type: self.write(": ") assign.type.accept(self) self.write(" = ") assign.value.accept(self) def visit_if(self, if_: If): self.write("if ") if_.condition.accept(self) self.write(":\n") self.indent() self.dispatch_write("\n", if_.true_body) self.dedent() if first_or_none(if_.false_body) is not None: self.write("else:\n") self.indent() self.dispatch_write("\n", if_.false_body) self.dedent() def visit_raise(self, raise_: Raise): self.write("raise ") raise_.expr.accept(self) def visit_functiondef(self, functiondef: FunctionDef): self.write("def ", functiondef.name, "(") self.dispatch_write(", ", functiondef.args) self.write(")") if functiondef.return_type: self.write(" -> ") functiondef.return_type.accept(self) self.write(":\n") self.indent() if functiondef.doc: self.write('r"""') for line in dedent(functiondef.doc).splitlines(): self.write(line) self.write("\n") self.write('"""\n\n') self.dispatch_write("\n", functiondef.body) self.dedent() def visit_classdef(self, classdef: ClassDef): self.write("class ", classdef.name) if any(classdef.bases): self.write("(") self.dispatch_write(", ", classdef.bases) self.write(")") self.write(":\n") self.indent() self.dispatch_write("\n\n", classdef.body) self.dedent() def visit_return(self, return_: Return): self.write("return ") return_.expr.accept(self) class DocCommentBuilder(Visitor): def __init__(self, width: int = 80): super().__init__() self.width = width def visit_functiondef(self, functiondef: FunctionDef): def wrap(text: str, initial_indent="", subsequent_indent=""): return TextWrapper( width=self.width, initial_indent=initial_indent, subsequent_indent=subsequent_indent, expand_tabs=False, replace_whitespace=False, fix_sentence_endings=False, break_long_words=False, break_on_hyphens=False, ).fill(text) argsdoc = "" for arg in functiondef.args: if arg.doc: argsdoc += wrap(f"{arg.name}: {arg.doc}", " " * 4, " " * 8) + "\n\n" if argsdoc: functiondef.doc += "\n\nArgs:\n" + argsdoc if functiondef.doc: functiondef.doc = functiondef.doc.strip()