Python ast.fix_missing_locations() Examples

def visit_Expr(self, node):
        newnode = ast.copy_location(ast.Expr(
                value = ast.Call(
                    func = ast.Attribute(
                        value = ast.Name(id='__swirlypy_recorder__',
                                    ctx=ast.Load()),
                                    attr="record",
                                    ctx=ast.Load()),
                                    args=[node.value],
                                    keywords=[],
                                    starargs=None,
                                    kwargs=None
                        )
                    ),
                node)
        ast.fix_missing_locations(newnode)
        return newnode 

def parse_entry_points_setup_cfg(visitor: Visitor) -> None:
    if not os.path.exists('setup.cfg'):
        return

    with visitor.file_ctx('setup.cfg', is_test=False):
        parser = configparser.ConfigParser()
        parser.read('setup.cfg')
        if 'options.entry_points' not in parser:
            return

        section = parser['options.entry_points']
        for k, v in section.items():
            for line in v.strip().splitlines():
                match = ENTRYPOINT_RE.match(line)
                if match:
                    node = ast.fix_missing_locations(ast.Str(match.group(1)))
                    visitor.read(match.group(1), node) 

def preprocess_method_body(source_code_obj):

    src = inspect.getsource(source_code_obj)

    ast_tree = ast.parse(src)

    visitor = RewriteDicts()
    ast_tree = visitor.visit(ast_tree)

    ast.fix_missing_locations(ast_tree)
    updated_code = compile(ast_tree, filename='<ast>', mode='exec')
    bc = Bytecode.from_code(updated_code)

    dlist = visitor.updated_dicts
    RewriteDicts.updated_dicts = []
    RewriteDicts.last_store_name = None

    block_code = get_code_block(bc)
    return block_code.arg, dlist 

def transform_ast(self, node):
        """Apply the AST transformations from self.ast_transformers
        
        Parameters
        ----------
        node : ast.Node
          The root node to be transformed. Typically called with the ast.Module
          produced by parsing user input.
        
        Returns
        -------
        An ast.Node corresponding to the node it was called with. Note that it
        may also modify the passed object, so don't rely on references to the
        original AST.
        """
        for transformer in self.ast_transformers:
            try:
                node = transformer.visit(node)
            except Exception:
                warn("AST transformer %r threw an error. It will be unregistered." % transformer)
                self.ast_transformers.remove(transformer)
        
        if self.ast_transformers:
            ast.fix_missing_locations(node)
        return node 

def updatePyAstVariableUsages(pyAst, variablesInScope, varsToTurnIntoCalls, isClassContext):
    #in the variables we've been handed, every member access chain 'x.y.z' has been
    #replaced with an actual variable lookup of the form 'x.y.z'. We need to perform
    #this same replacement in the actual python source code we're running, since
    #we may not actually have enough information to recover 'x'

    replacements = {}
    for possible_replacement in variablesInScope:
        if '.' in possible_replacement:
            replacements[tuple(possible_replacement.split('.'))] = possible_replacement

    #we need to deepcopy the AST since the collapser modifies the AST, and this is just a 
    #slice of a cached tree.
    pyAst = ast.fix_missing_locations(copy.deepcopy(pyAst))

    if varsToTurnIntoCalls:
        pyAst = PyAstFreeVariableAnalyses.replaceUsesWithCalls(pyAst, set(varsToTurnIntoCalls), isClassContext)

    pyAst = PyAstFreeVariableAnalyses.collapseFreeVariableMemberAccessChains(
        pyAst,
        replacements,
        isClassContext=isClassContext
        )

    return ast.fix_missing_locations(pyAst) 

def visit_BinOp(self, node):
        'If node is a constant expression, replace it with its evaluated value'
        if node in self.constexpr:
            # evaluation
            fake_node = ast.Expression(ast.BinOp(node, ast.Mod(),
                                                 ast.Num(self.mod)))
            ast.fix_missing_locations(fake_node)
            code = compile(fake_node, '<constant folding>', 'eval')
            obj_env = globals().copy()
            exec code in obj_env
            value = eval(code, obj_env)

            new_node = ast.Num(value)
            return new_node
        else:
            return self.generic_visit(node) 

def visit_UnaryOp(self, node):
        'Same idea as visit_BinOp'
        if node in self.constexpr:
            # evaluation
            fake_node = ast.Expression(ast.BinOp(node, ast.Mod(),
                                                 ast.Num(self.mod)))
            ast.fix_missing_locations(fake_node)
            code = compile(fake_node, '<constant folding>', 'eval')
            obj_env = globals().copy()
            exec code in obj_env

            value = eval(code, obj_env)
            new_node = ast.Num(value)
            return new_node
        else:
            return self.generic_visit(node) 

def test_invalid_string(self):
        m = ast.Module([ast.Expr(ast.Str(42))])
        ast.fix_missing_locations(m)
        with self.assertRaises(TypeError) as cm:
            compile(m, "<test>", "exec")
        self.assertIn("string must be of type str", str(cm.exception)) 

def visit_stmt(self, node):
        # type: (ast.stmt) -> ast.With
        """
        Every statement in the original code becomes:

        with _treetrace_hidden_with_stmt(_tree_index):
            <statement>

        where the _treetrace_hidden_with_stmt function is the the corresponding method with the
        TreeTracerBase and traced_file arguments already filled in (see _trace_methods_dict)
        """
        context_expr = self._create_simple_marker_call(
            super(_NodeVisitor, self).generic_visit(node),
            TreeTracerBase._treetrace_hidden_with_stmt)

        if PY3:
            wrapped = ast.With(
                items=[ast.withitem(context_expr=context_expr)],
                body=[node],
            )
        else:
            wrapped = ast.With(
                context_expr=context_expr,
                body=[node],
            )
        ast.copy_location(wrapped, node)
        ast.fix_missing_locations(wrapped)
        return wrapped 

def wrap_code(code: str, args: str = '') -> ast.Module:
    """
    Compiles Python code into an async function or generator,
    and automatically adds return if the function body is a single evaluation.
    Also adds inline import expression support.
    """

    user_code = import_expression.parse(code, mode='exec')
    mod = import_expression.parse(CORO_CODE.format(args), mode='exec')

    definition = mod.body[-1]  # async def ...:
    assert isinstance(definition, ast.AsyncFunctionDef)

    try_block = definition.body[-1]  # try:
    assert isinstance(try_block, ast.Try)

    try_block.body.extend(user_code.body)

    ast.fix_missing_locations(mod)

    KeywordTransformer().generic_visit(try_block)

    last_expr = try_block.body[-1]

    # if the last part isn't an expression, ignore it
    if not isinstance(last_expr, ast.Expr):
        return mod

    # if the last expression is not a yield
    if not isinstance(last_expr.value, ast.Yield):
        # copy the value of the expression into a yield
        yield_stmt = ast.Yield(last_expr.value)
        ast.copy_location(yield_stmt, last_expr)
        # place the yield into its own expression
        yield_expr = ast.Expr(yield_stmt)
        ast.copy_location(yield_expr, last_expr)

        # place the yield where the original expression was
        try_block.body[-1] = yield_expr

    return mod 

def test_fix_missing_locations(self):
        src = ast.parse('write("spam")')
        src.body.append(ast.Expr(ast.Call(ast.Name('spam', ast.Load()),
                                          [ast.Str('eggs')], [], None, None)))
        self.assertEqual(src, ast.fix_missing_locations(src))
        self.assertEqual(ast.dump(src, include_attributes=True),
            "Module(body=[Expr(value=Call(func=Name(id='write', ctx=Load(), "
            "lineno=1, col_offset=0), args=[Str(s='spam', lineno=1, "
            "col_offset=6)], keywords=[], starargs=None, kwargs=None, "
            "lineno=1, col_offset=0), lineno=1, col_offset=0), "
            "Expr(value=Call(func=Name(id='spam', ctx=Load(), lineno=1, "
            "col_offset=0), args=[Str(s='eggs', lineno=1, col_offset=0)], "
            "keywords=[], starargs=None, kwargs=None, lineno=1, "
            "col_offset=0), lineno=1, col_offset=0)])"
        ) 

def test_invalid_identitifer(self):
        m = ast.Module([ast.Expr(ast.Name(42, ast.Load()))])
        ast.fix_missing_locations(m)
        with self.assertRaises(TypeError) as cm:
            compile(m, "<test>", "exec")
        self.assertIn("identifier must be of type str", str(cm.exception)) 

def _execute_comprehension(self, node: Union[ast.ListComp, ast.SetComp, ast.GeneratorExp, ast.DictComp]) -> Any:
        """Compile the generator or comprehension from the node and execute the compiled code."""
        args = [ast.arg(arg=name) for name in sorted(self._name_to_value.keys())]

        if platform.python_version_tuple() < ('3', ):
            raise NotImplementedError("Python versions below not supported, got: {}".format(platform.python_version()))

        if platform.python_version_tuple() < ('3', '8'):
            func_def_node = ast.FunctionDef(
                name="generator_expr",
                args=ast.arguments(args=args, kwonlyargs=[], kw_defaults=[], defaults=[]),
                decorator_list=[],
                body=[ast.Return(node)])

            module_node = ast.Module(body=[func_def_node])
        else:
            func_def_node = ast.FunctionDef(
                name="generator_expr",
                args=ast.arguments(args=args, posonlyargs=[], kwonlyargs=[], kw_defaults=[], defaults=[]),
                decorator_list=[],
                body=[ast.Return(node)])

            module_node = ast.Module(body=[func_def_node], type_ignores=[])

        ast.fix_missing_locations(module_node)

        code = compile(source=module_node, filename='<ast>', mode='exec')

        module_locals = {}  # type: Dict[str, Any]
        module_globals = {}  # type: Dict[str, Any]
        exec(code, module_globals, module_locals)  # pylint: disable=exec-used

        generator_expr_func = module_locals["generator_expr"]

        return generator_expr_func(**self._name_to_value) 

def test_fix_missing_locations(self):
        src = ast.parse('write("spam")')
        src.body.append(ast.Expr(ast.Call(ast.Name('spam', ast.Load()),
                                          [ast.Str('eggs')], [])))
        self.assertEqual(src, ast.fix_missing_locations(src))
        self.assertEqual(ast.dump(src, include_attributes=True),
            "Module(body=[Expr(value=Call(func=Name(id='write', ctx=Load(), "
            "lineno=1, col_offset=0), args=[Str(s='spam', lineno=1, "
            "col_offset=6)], keywords=[], "
            "lineno=1, col_offset=0), lineno=1, col_offset=0), "
            "Expr(value=Call(func=Name(id='spam', ctx=Load(), lineno=1, "
            "col_offset=0), args=[Str(s='eggs', lineno=1, col_offset=0)], "
            "keywords=[], lineno=1, "
            "col_offset=0), lineno=1, col_offset=0)])"
        ) 

def mod(self, mod, msg=None, mode="exec", *, exc=ValueError):
        mod.lineno = mod.col_offset = 0
        ast.fix_missing_locations(mod)
        with self.assertRaises(exc) as cm:
            compile(mod, "<test>", mode)
        if msg is not None:
            self.assertIn(msg, str(cm.exception)) 

def run_call(args, node, process, get_func, **kwargs):
    # Get function expression
    if isinstance(node, ast.FunctionDef):  # function name
        func_expr = ast.Name(id=node.name, ctx=ast.Load())
    elif isinstance(node, ast.Lambda):  # lambda body expr
        func_expr = node
    else:
        raise InstructorError.from_message(
            "Only function definition or lambda may be called"
        )

    ast.fix_missing_locations(func_expr)
    return get_func(process=process, tree=func_expr, call=args, **kwargs) 

def assign_from_ast(node, expr):
    """
    Creates code to assign name (or tuple of names) node from expr

    This is useful for recreating destructuring assignment behavior, like
    a, *b = [1,2,3].
    """
    if isinstance(expr, str):
        expr = ast.Name(id=expr, ctx=ast.Load())
    mod = ast.Module([ast.Assign(targets=[node], value=expr)])
    ast.fix_missing_locations(mod)
    return compile(mod, "<assignment_script>", "exec") 

def ConvertFuncDecl(self, func_desc):
        # assert(func_desc.func_ast is not None)
        # ast.fix_missing_locations(func_desc.func_ast)
        # func_desc.Calculate_StackSize(self.global_num)

        # build dynamic stack first
        # self.tokenizer.build_function_stack(func_desc.stack_size, func_desc.func_ast)

        # load the argument which passed by caller into stackscope
        CodeGenVisitor = Visitor_Of_FunCodeGen(self, func_desc)
        CodeGenVisitor.visit(func_desc.func_ast)
        # self.tokenizer.dump_all_vm_token()

    # bring all import func into funscope. Include compile builtinl 

def visit_expr(self, node):
        # type: (ast.expr) -> ast.Call
        """
        each expression e gets wrapped like this:
            _treetrace_hidden_after_expr(_treetrace_hidden_before_expr(_tree_index), e)

        where the _treetrace_* functions are the corresponding methods with the
        TreeTracerBase and traced_file arguments already filled in (see _trace_methods_dict)
        """

        before_marker = self._create_simple_marker_call(node, TreeTracerBase._treetrace_hidden_before_expr)
        ast.copy_location(before_marker, node)

        after_marker = ast.Call(
            func=ast.Name(id=TreeTracerBase._treetrace_hidden_after_expr.__name__,
                          ctx=ast.Load()),
            args=[
                before_marker,
                super(_NodeVisitor, self).generic_visit(node),
            ],
            keywords=[],
        )
        ast.copy_location(after_marker, node)
        ast.fix_missing_locations(after_marker)

        return after_marker 

def matching_nodes(self, exprs):
        for i, expr in enumerate(exprs):
            setter = get_setter(expr)
            replacement = ast.BinOp(
                left=expr,
                op=ast.Pow(),
                right=ast.Str(s=sentinel),
            )
            ast.fix_missing_locations(replacement)
            setter(replacement)
            try:
                instructions = self.compile_instructions()
            except SyntaxError:
                continue
            finally:
                setter(expr)
            indices = [
                i
                for i, instruction in enumerate(instructions)
                if instruction.argval == sentinel
            ]
            if not indices:
                continue
            arg_index = only(indices) - 1
            while instructions[arg_index].opname == 'EXTENDED_ARG':
                arg_index -= 1

            if instructions[arg_index].offset == self.frame.f_lasti:
                yield expr 

def _rewrite_function(f):
    """
    Rewrite a function so that any if/else branches are intercepted and their behavior can be
    overridden. This is accomplished using 3 steps:

    1. Get the source of the function and then rewrite the AST using _IfTransformer
    2. Do some small fixups to the tree to make sure
        a) the function doesn't have the same name, and
        b) the decorator isn't called recursively on the transformed function as well
    3. Bring the variables from the call site back into scope

    :param f: Function to rewrite
    :return: Rewritten function
    """
    source = inspect.getsource(f)
    tree = ast.parse(source)
    _IfTransformer().visit(tree)

    ast.fix_missing_locations(tree)
    tree.body[0].name = f.__name__ + "_patched"
    tree.body[0].decorator_list = []

    compiled = compile(tree, filename="<ast>", mode="exec")
    # The first f_back here gets to the body of magicquil() and the second f_back gets to the
    # user's call site which is what we want. If we didn't add these manually to the globals it
    # wouldn't be possible to call other @magicquil functions from within a @magicquil function.
    prev_globals = inspect.currentframe().f_back.f_back.f_globals
    # For reasons I don't quite understand it's critical to add locals() here otherwise the
    # function will disappear and we won't be able to return it below
    exec(compiled, {**prev_globals, **globals()}, locals())
    return locals()[f.__name__ + "_patched"] 

def test_invalid_identitifer(self):
        m = ast.Module([ast.Expr(ast.Name(u"x", ast.Load()))])
        ast.fix_missing_locations(m)
        with self.assertRaises(TypeError) as cm:
            compile(m, "<test>", "exec")
        self.assertIn("identifier must be of type str", str(cm.exception)) 

def test_compile_manual(self):
        # check that expressions which are built manually compile
        for test in eval_tests:
            a = ast.parse(test, "<unknown>", mode="eval")
            b = ast.fix_missing_locations(eval(ast.dump(a, annotate_fields=False), vars(ast)))
            compile(b, "<unknown>", mode="eval")
        for test in exec_tests:
            a = ast.parse(test, "<unknown>", mode="exec")
            b = ast.fix_missing_locations(eval(ast.dump(a, annotate_fields=False), vars(ast)))
            compile(b, "<unknown>", mode="exec") 

def mod(self, mod, msg=None, mode="exec", *, exc=ValueError):
        mod.lineno = mod.col_offset = 0
        ast.fix_missing_locations(mod)
        with self.assertRaises(exc) as cm:
            compile(mod, "<test>", mode)
        if msg is not None:
            self.assertIn(msg, str(cm.exception)) 

def test_fix_missing_locations(self):
        src = ast.parse('write("spam")')
        src.body.append(ast.Expr(ast.Call(ast.Name('spam', ast.Load()),
                                          [ast.Str('eggs')], [], None, None)))
        self.assertEqual(src, ast.fix_missing_locations(src))
        self.assertEqual(ast.dump(src, include_attributes=True),
            "Module(body=[Expr(value=Call(func=Name(id='write', ctx=Load(), "
            "lineno=1, col_offset=0), args=[Str(s='spam', lineno=1, "
            "col_offset=6)], keywords=[], starargs=None, kwargs=None, "
            "lineno=1, col_offset=0), lineno=1, col_offset=0), "
            "Expr(value=Call(func=Name(id='spam', ctx=Load(), lineno=1, "
            "col_offset=0), args=[Str(s='eggs', lineno=1, col_offset=0)], "
            "keywords=[], starargs=None, kwargs=None, lineno=1, "
            "col_offset=0), lineno=1, col_offset=0)])"
        ) 

def test_invalid_string(self):
        m = ast.Module([ast.Expr(ast.Str(42))])
        ast.fix_missing_locations(m)
        with self.assertRaises(TypeError) as cm:
            compile(m, "<test>", "exec")
        self.assertIn("string must be of type str", str(cm.exception)) 

def test_invalid_identitifer(self):
        m = ast.Module([ast.Expr(ast.Name(42, ast.Load()))])
        ast.fix_missing_locations(m)
        with self.assertRaises(TypeError) as cm:
            compile(m, "<test>", "exec")
        self.assertIn("identifier must be of type str", str(cm.exception)) 

def visit_Assign(self, node):
        """Catch all the globals."""
        self.generic_visit(node)

        if node.col_offset == 0:
            mnode = ast.parse("")
            mnode.body = [node]
            mnode = ast.fix_missing_locations(mnode)
            code = compile(mnode, "<ast>", "exec")
            try:
                exec(code, self.globals_)
            except Exception:
                pass
        self.globals_.pop("__builtins__", None)
        self.globals_.pop("builtins", None) 

def insert_returns(body: List[ast.AST]) -> None:
    if isinstance(body[-1], ast.Expr):
        body[-1] = ast.Return(body[-1].value)
        ast.fix_missing_locations(body[-1])
    elif isinstance(body[-1], ast.If):
        insert_returns(body[-1].body)
        insert_returns(body[-1].orelse)
    elif isinstance(body[-1], (ast.With, ast.AsyncWith)):
        insert_returns(body[-1].body) 

def make_call_node(node: ast.expr, func_name: str) -> ast.Expr:
    msg = get_assertion_msg(node)
    call = ast.Call(
        func=ast.Name(id=func_name, ctx=ast.Load()),
        args=[node.test.left, node.test.comparators[0], ast.Str(s=msg)],
        keywords=[],
    )
    new_node = ast.Expr(value=call)
    ast.copy_location(new_node, node)
    ast.fix_missing_locations(new_node)
    return new_node