Пример #1
0
    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,
        )
Пример #2
0
    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),
        )
Пример #3
0
    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)