def leave_Yield(self, node: cst.Yield, updated_node: cst.Yield) -> Union[cst.Await, cst.Yield]: if not self.in_coroutine(self.coroutine_stack): return updated_node if not isinstance(updated_node.value, cst.BaseExpression): return updated_node if isinstance(updated_node.value, (cst.List, cst.ListComp)): self.required_imports.add("asyncio") expression = self.pluck_asyncio_gather_expression_from_yield_list_or_list_comp( updated_node) elif m.matches( updated_node, m.Yield(value=((m.Dict() | m.DictComp())) | m.Call(func=m.Name("dict"))), ): raise TransformError( "Yielding a dict of futures (https://www.tornadoweb.org/en/branch3.2/releases/v3.2.0.html#tornado-gen) added in tornado 3.2 is unsupported by the codemod. This file has not been modified. Manually update to supported syntax before running again." ) else: expression = updated_node.value return cst.Await( expression=expression, whitespace_after_await=updated_node.whitespace_after_yield, lpar=updated_node.lpar, rpar=updated_node.rpar, )
def leave_FunctionDef( self, original_node: cst.FunctionDef, updated_node: cst.FunctionDef ) -> Union[cst.BaseStatement, cst.RemovalSentinel]: modified_defaults: List = [] mutable_args: List[Tuple[cst.Name, Union[cst.List, cst.Dict]]] = [] for param in updated_node.params.params: if not m.matches(param, m.Param(default=m.OneOf(m.List(), m.Dict()))): modified_defaults.append(param) continue # This line here is just for type checkers peace of mind, # since it cannot reason about variables from matchers result. if not isinstance(param.default, (cst.List, cst.Dict)): continue mutable_args.append((param.name, param.default)) modified_defaults.append( param.with_changes(default=cst.Name("None"), )) if not mutable_args: return original_node modified_params: cst.Parameters = updated_node.params.with_changes( params=modified_defaults) initializations: List[Union[ cst.SimpleStatementLine, cst.BaseCompoundStatement]] = [ # We use generation by template here since construction of the # resulting 'if' can be burdensome due to many nested objects # involved. Additional line is attached so that we may control # exact spacing between generated statements. parse_template_statement( DEFAULT_INIT_TEMPLATE, config=self.module_config, arg=arg, init=init).with_changes(leading_lines=[EMPTY_LINE]) for arg, init in mutable_args ] # Docstring should always go right after the function definition, # so we take special care to insert our initializations after the # last docstring found. docstrings = takewhile(is_docstring, updated_node.body.body) function_code = dropwhile(is_docstring, updated_node.body.body) # It is not possible to insert empty line after the statement line, # because whitespace is owned by the next statement after it. stmt_with_empty_line = next(function_code).with_changes( leading_lines=[EMPTY_LINE]) modified_body = ( *docstrings, *initializations, stmt_with_empty_line, *function_code, ) return updated_node.with_changes( params=modified_params, body=updated_node.body.with_changes(body=modified_body), )
def obf_universal(self, node: cst.CSTNode, *types): if m.matches(node, m.Name()): types = ('a', 'ca', 'v', 'cv') if not types else types node = cst.ensure_type(node, cst.Name) if self.can_rename(node.value, *types): node = self.get_new_cst_name(node) elif m.matches(node, m.NameItem()): node = cst.ensure_type(node, cst.NameItem) node = node.with_changes(name=self.obf_universal(node.name)) elif m.matches(node, m.Call()): node = cst.ensure_type(node, cst.Call) if self.change_methods or self.change_functions: node = self.new_obf_function_name(node) if self.change_arguments or self.change_method_arguments: node = self.obf_function_args(node) elif m.matches(node, m.Attribute()): node = cst.ensure_type(node, cst.Attribute) value = node.value attr = node.attr self.obf_universal(value) self.obf_universal(attr) elif m.matches(node, m.AssignTarget()): node = cst.ensure_type(node, cst.AssignTarget) node = node.with_changes(target=self.obf_universal(node.target)) elif m.matches(node, m.List() | m.Tuple()): node = cst.ensure_type(node, cst.List) if m.matches( node, m.List()) else cst.ensure_type(node, cst.Tuple) new_elements = [] for el in node.elements: new_elements.append(self.obf_universal(el)) node = node.with_changes(elements=new_elements) elif m.matches(node, m.Subscript()): node = cst.ensure_type(node, cst.Subscript) new_slice = [] for el in node.slice: new_slice.append( el.with_changes(slice=self.obf_slice(el.slice))) node = node.with_changes(slice=new_slice) node = node.with_changes(value=self.obf_universal(node.value)) elif m.matches(node, m.Element()): node = cst.ensure_type(node, cst.Element) node = node.with_changes(value=self.obf_universal(node.value)) elif m.matches(node, m.Dict()): node = cst.ensure_type(node, cst.Dict) new_elements = [] for el in node.elements: new_elements.append(self.obf_universal(el)) node = node.with_changes(elements=new_elements) elif m.matches(node, m.DictElement()): node = cst.ensure_type(node, cst.DictElement) new_key = self.obf_universal(node.key) new_val = self.obf_universal(node.value) node = node.with_changes(key=new_key, value=new_val) elif m.matches(node, m.StarredDictElement()): node = cst.ensure_type(node, cst.StarredDictElement) node = node.with_changes(value=self.obf_universal(node.value)) elif m.matches(node, m.If() | m.While()): node = cst.ensure_type(node, cst.IfExp) if m.matches( node, cst.If | cst.IfExp) else cst.ensure_type(node, cst.While) node = node.with_changes(test=self.obf_universal(node.test)) elif m.matches(node, m.IfExp()): node = cst.ensure_type(node, cst.IfExp) node = node.with_changes(body=self.obf_universal(node.body)) node = node.with_changes(test=self.obf_universal(node.test)) node = node.with_changes(orelse=self.obf_universal(node.orelse)) elif m.matches(node, m.Comparison()): node = cst.ensure_type(node, cst.Comparison) new_compars = [] for target in node.comparisons: new_compars.append(self.obf_universal(target)) node = node.with_changes(left=self.obf_universal(node.left)) node = node.with_changes(comparisons=new_compars) elif m.matches(node, m.ComparisonTarget()): node = cst.ensure_type(node, cst.ComparisonTarget) node = node.with_changes( comparator=self.obf_universal(node.comparator)) elif m.matches(node, m.FormattedString()): node = cst.ensure_type(node, cst.FormattedString) new_parts = [] for part in node.parts: new_parts.append(self.obf_universal(part)) node = node.with_changes(parts=new_parts) elif m.matches(node, m.FormattedStringExpression()): node = cst.ensure_type(node, cst.FormattedStringExpression) node = node.with_changes( expression=self.obf_universal(node.expression)) elif m.matches(node, m.BinaryOperation() | m.BooleanOperation()): node = cst.ensure_type(node, cst.BinaryOperation) if m.matches( node, m.BinaryOperation()) else cst.ensure_type( node, cst.BooleanOperation) node = node.with_changes(left=self.obf_universal(node.left), right=self.obf_universal(node.right)) elif m.matches(node, m.UnaryOperation()): node = cst.ensure_type(node, cst.UnaryOperation) node = node.with_changes( expression=self.obf_universal(node.expression)) elif m.matches(node, m.ListComp()): node = cst.ensure_type(node, cst.ListComp) node = node.with_changes(elt=self.obf_universal(node.elt)) node = node.with_changes(for_in=self.obf_universal(node.for_in)) elif m.matches(node, m.DictComp()): node = cst.ensure_type(node, cst.DictComp) node = node.with_changes(key=self.obf_universal(node.key)) node = node.with_changes(value=self.obf_universal(node.value)) node = node.with_changes(for_in=self.obf_universal(node.for_in)) elif m.matches(node, m.CompFor()): node = cst.ensure_type(node, cst.CompFor) new_ifs = [] node = node.with_changes(target=self.obf_universal(node.target)) node = node.with_changes(iter=self.obf_universal(node.iter)) for el in node.ifs: new_ifs.append(self.obf_universal(el)) node = node.with_changes(ifs=new_ifs) elif m.matches(node, m.CompIf()): node = cst.ensure_type(node, cst.CompIf) node = node.with_changes(test=self.obf_universal(node.test)) elif m.matches(node, m.Integer() | m.Float() | m.SimpleString()): pass else: pass # print(node) return node
def visit_Call(self, node: cst.Call) -> None: result = m.extract( node, m.Call( func=m.Attribute(value=m.Name("self"), attr=m.Name("assertTrue")), args=[ m.DoNotCare(), m.Arg(value=m.SaveMatchedNode( m.OneOf( m.Integer(), m.Float(), m.Imaginary(), m.Tuple(), m.List(), m.Set(), m.Dict(), m.Name("None"), m.Name("True"), m.Name("False"), ), "second", )), ], ), ) if result: second_arg = result["second"] if isinstance(second_arg, Sequence): second_arg = second_arg[0] if m.matches(second_arg, m.Name("True")): new_call = node.with_changes(args=[ node.args[0].with_changes(comma=cst.MaybeSentinel.DEFAULT) ], ) elif m.matches(second_arg, m.Name("None")): new_call = node.with_changes( func=node.func.with_deep_changes( old_node=cst.ensure_type(node.func, cst.Attribute).attr, value="assertIsNone", ), args=[ node.args[0].with_changes( comma=cst.MaybeSentinel.DEFAULT) ], ) elif m.matches(second_arg, m.Name("False")): new_call = node.with_changes( func=node.func.with_deep_changes( old_node=cst.ensure_type(node.func, cst.Attribute).attr, value="assertFalse", ), args=[ node.args[0].with_changes( comma=cst.MaybeSentinel.DEFAULT) ], ) else: new_call = node.with_deep_changes( old_node=cst.ensure_type(node.func, cst.Attribute).attr, value="assertEqual", ) self.report(node, replacement=new_call)