Python ast.Assign() Examples

def version():
    path = 'pypika/__init__.py'
    with open(path, 'r') as file:
        t = compile(file.read(), path, 'exec', ast.PyCF_ONLY_AST)
        for node in (n for n in t.body if isinstance(n, ast.Assign)):
            if len(node.targets) == 1:
                name = node.targets[0]
                if isinstance(name, ast.Name) and \
                      name.id in ('__version__', '__version_info__', 'VERSION'):
                    v = node.value
                    if isinstance(v, ast.Str):
                        return v.s

                    if isinstance(v, ast.Tuple):
                        r = []
                        for e in v.elts:
                            if isinstance(e, ast.Str):
                                r.append(e.s)
                            elif isinstance(e, ast.Num):
                                r.append(str(e.n))
                        return '.'.join(r) 

def pop_format_context(self, expl_expr):
        """Format the %-formatted string with current format context.

        The expl_expr should be an ast.Str instance constructed from
        the %-placeholders created by .explanation_param().  This will
        add the required code to format said string to .expl_stmts and
        return the ast.Name instance of the formatted string.

        """
        current = self.stack.pop()
        if self.stack:
            self.explanation_specifiers = self.stack[-1]
        keys = [ast.Str(key) for key in current.keys()]
        format_dict = ast.Dict(keys, list(current.values()))
        form = ast.BinOp(expl_expr, ast.Mod(), format_dict)
        name = "@py_format" + str(next(self.variable_counter))
        if self.enable_assertion_pass_hook:
            self.format_variables.append(name)
        self.expl_stmts.append(ast.Assign([ast.Name(name, ast.Store())], form))
        return ast.Name(name, ast.Load()) 

def visit_FunctionDef(self, node):
        self.refStack.push()
        for n in node.body:
            self.visit(n)
        self.tree.kind = _kind.SIMPLE_ALWAYS_COMB
        for n in node.body:
            if isinstance(n, ast.Expr) and isinstance(n.value, ast.Str):
                continue  # skip doc strings
            if isinstance(n, ast.Assign) and \
               isinstance(n.targets[0], ast.Attribute) and \
               self.getKind(n.targets[0].value) != _kind.REG:
                pass
            else:
                self.tree.kind = _kind.ALWAYS_COMB
                return
        # rom access is expanded into a case statement in addition
        # to any always_comb that contains a list of signals
        # if self.tree.hasRom or self.tree.hasLos:
        if self.tree.hasRom:
            self.tree.kind = _kind.ALWAYS_COMB
        self.refStack.pop() 

def expr_stmt_rewrite(lhs, ann, aug, aug_exp, rhs: t.Optional[list]):
    if rhs:
        as_store(lhs)
        *init, end = rhs
        for each in init:
            as_store(each)
        return ast.Assign([lhs, *init], end)

    if ann:
        as_store(lhs)
        anno, value = ann
        return ast.AnnAssign(lhs, anno, value, 1)

    if aug_exp:
        as_store(lhs)
        return ast.AugAssign(lhs, aug(), aug_exp)

    # NO AS STORE HERE!
    return ast.Expr(lhs) 

def get_config_comments(func):
    filename = inspect.getfile(func)
    func_body, line_offset = get_function_body(func)
    body_source = dedent_function_body(func_body)
    body_code = compile(body_source, filename, "exec", ast.PyCF_ONLY_AST)
    body_lines = body_source.split("\n")

    variables = {"seed": "the random seed for this experiment"}

    for ast_root in body_code.body:
        for ast_entry in [ast_root] + list(ast.iter_child_nodes(ast_root)):
            if isinstance(ast_entry, ast.Assign):
                # we found an assignment statement
                # go through all targets of the assignment
                # usually a single entry, but can be more for statements like:
                # a = b = 5
                for t in ast_entry.targets:
                    add_doc(t, variables, body_lines)

    return variables 

def visit_Assign(self, node: ast.Assign) -> None:
        for target in node.targets:
            if isinstance(target, ast.Name):
                self.define(target.id, node)

        if (
                len(node.targets) == 1 and
                isinstance(node.targets[0], ast.Name) and
                node.targets[0].id == '__all__' and
                isinstance(node.value, (ast.Tuple, ast.List))
        ):
            for elt in node.value.elts:
                if isinstance(elt, ast.Str):
                    self.read(elt.s, elt)

        self.generic_visit(node)

    # TODO: AnnAssign 

def pop_format_context(self, expl_expr):
        """Format the %-formatted string with current format context.

        The expl_expr should be an ast.Str instance constructed from
        the %-placeholders created by .explanation_param().  This will
        add the required code to format said string to .expl_stmts and
        return the ast.Name instance of the formatted string.

        """
        current = self.stack.pop()
        if self.stack:
            self.explanation_specifiers = self.stack[-1]
        keys = [ast.Str(key) for key in current.keys()]
        format_dict = ast.Dict(keys, list(current.values()))
        form = ast.BinOp(expl_expr, ast.Mod(), format_dict)
        name = "@py_format" + str(next(self.variable_counter))
        if self.enable_assertion_pass_hook:
            self.format_variables.append(name)
        self.expl_stmts.append(ast.Assign([ast.Name(name, ast.Store())], form))
        return ast.Name(name, ast.Load()) 

def test_with_src_module_file():
    poetry = Factory().create_poetry(project("source_file"))

    builder = SdistBuilder(poetry, NullEnv(), NullIO())

    # Check setup.py
    setup = builder.build_setup()
    setup_ast = ast.parse(setup)

    setup_ast.body = [n for n in setup_ast.body if isinstance(n, ast.Assign)]
    ns = {}
    exec(compile(setup_ast, filename="setup.py", mode="exec"), ns)
    assert ns["package_dir"] == {"": "src"}
    assert ns["modules"] == ["module_src"]

    builder.build()

    sdist = fixtures_dir / "source_file" / "dist" / "module-src-0.1.tar.gz"

    assert sdist.exists()

    with tarfile.open(str(sdist), "r") as tar:
        assert "module-src-0.1/src/module_src.py" in tar.getnames() 

def test_with_src_module_dir():
    poetry = Factory().create_poetry(project("source_package"))

    builder = SdistBuilder(poetry, NullEnv(), NullIO())

    # Check setup.py
    setup = builder.build_setup()
    setup_ast = ast.parse(setup)

    setup_ast.body = [n for n in setup_ast.body if isinstance(n, ast.Assign)]
    ns = {}
    exec(compile(setup_ast, filename="setup.py", mode="exec"), ns)
    assert ns["package_dir"] == {"": "src"}
    assert ns["packages"] == ["package_src"]

    builder.build()

    sdist = fixtures_dir / "source_package" / "dist" / "package-src-0.1.tar.gz"

    assert sdist.exists()

    with tarfile.open(str(sdist), "r") as tar:
        assert "package-src-0.1/src/package_src/__init__.py" in tar.getnames()
        assert "package-src-0.1/src/package_src/module.py" in tar.getnames() 

def _find_variable_in_body(
        self, body, name
    ):  # type: (Iterable[Any], str) -> Optional[Any]
        found = None
        for elem in body:
            if found:
                break

            if not isinstance(elem, ast.Assign):
                continue

            for target in elem.targets:
                if not isinstance(target, ast.Name):
                    continue

                if target.id == name:
                    return elem.value 

def visit_Return(self, return_):
        """Convert returns into assignment/exception pairs

        Since the body of this function will be in the global namespace we
        can't have any returns. An acceptable alternative is to set a variable
        called 'RETURN' and then immediately raise an exception.

        >>> self = NamespacePromoter(buffer='foo')
        >>> code = '''
        ...
        ... return 5
        ...
        ... '''
        >>> tree = ast.parse(code)
        >>> return_, = tree.body

        """
        nodes = [
            ast.Assign(targets=[ast.Name(id='RETURN', ctx=ast.Store())], value=return_.value, lineno=return_.lineno),
            ast.Raise(exc=ast.Call(func=ast.Name(id='Exception', ctx=ast.Load()), args=[ast.Str(s='return')], keywords=[]), cause=None),
        ]
        return nodes 

def _fine_property_definition(self, property_name):
        """Find the lines in the source code that contain this property's name and definition.

        This function can find both attribute assignments as well as methods/functions.

        Args:
            property_name (str): the name of the property to look up in the template definition

        Returns:
            tuple: line numbers for the start and end of the attribute definition
        """
        for node in ast.walk(ast.parse(self._source)):
            if isinstance(node, ast.Assign) and node.targets[0].id == property_name:
                return node.targets[0].lineno - 1, self._get_node_line_end(node)
            elif isinstance(node, ast.FunctionDef) and node.name == property_name:
                return node.lineno - 1, self._get_node_line_end(node)
        raise ValueError('The requested node could not be found.') 

def pop_format_context(self, expl_expr: ast.expr) -> ast.Name:
        """Format the %-formatted string with current format context.

        The expl_expr should be an str ast.expr instance constructed from
        the %-placeholders created by .explanation_param().  This will
        add the required code to format said string to .expl_stmts and
        return the ast.Name instance of the formatted string.

        """
        current = self.stack.pop()
        if self.stack:
            self.explanation_specifiers = self.stack[-1]
        keys = [ast.Str(key) for key in current.keys()]
        format_dict = ast.Dict(keys, list(current.values()))
        form = ast.BinOp(expl_expr, ast.Mod(), format_dict)
        name = "@py_format" + str(next(self.variable_counter))
        if self.enable_assertion_pass_hook:
            self.format_variables.append(name)
        self.expl_stmts.append(ast.Assign([ast.Name(name, ast.Store())], form))
        return ast.Name(name, ast.Load()) 

def _parse_setup_params(setup_py: Path) -> Dict[str, Any]:
    with setup_py.open("r", encoding="utf8") as sp:
        setup_tree = ast.parse(sp.read())

    LOG.debug(f"AST visiting {setup_py}")
    for node in ast.walk(setup_tree):
        if isinstance(node, ast.Assign):
            for target in node.targets:
                target_id = getattr(target, "id", None)
                if not target_id:
                    continue

                if target_id == "ptr_params":
                    LOG.debug(f"Found ptr_params in {setup_py}")
                    return dict(ast.literal_eval(node.value))
    return {} 

def read_tags(filename):
    """Reads values of "magic tags" defined in the given Python file.

    :param filename: Python filename to read the tags from
    :return: Dictionary of tags
    """
    with open(filename) as f:
        ast_tree = ast.parse(f.read(), filename)

    res = {}
    for node in ast.walk(ast_tree):
        if type(node) is not ast.Assign:
            continue

        target = node.targets[0]
        if type(target) is not ast.Name:
            continue
        if not (target.id.startswith('__') and target.id.endswith('__')):
            continue

        name = target.id[2:-2]
        res[name] = ast.literal_eval(node.value)

    return res 

def assign_multi_target(self, node, right_hand_side_variables):
        new_assignment_nodes = list()

        for target in node.targets:
            label = LabelVisitor()
            label.visit(target)
            left_hand_side = label.result
            label.result += ' = '
            label.visit(node.value)
            new_assignment_nodes.append(self.append_node(AssignmentNode(
                label.result,
                left_hand_side,
                ast.Assign(target, node.value),
                right_hand_side_variables,
                line_number=node.lineno,
                path=self.filenames[-1]
            )))

        connect_nodes(new_assignment_nodes)
        return ControlFlowNode(new_assignment_nodes[0], [new_assignment_nodes[-1]], [])  # return the last added node 

def get_version(fname):
    """Get the version info from the mpld3 package without importing it"""
    import ast

    with open(fname) as init_file:
        module = ast.parse(init_file.read())

    version = (ast.literal_eval(node.value) for node in ast.walk(module)
               if isinstance(node, ast.Assign)
               and node.targets[0].id == "__version__")
    try:
        return next(version)
    except StopIteration:
        raise ValueError("version could not be located")


#cffi.verifier.cleanup_tmpdir() 

def _find_has_import(self):
        for child in self.ast.body:
            found_try_except_import = False
            found_has = False
            if isinstance(child, ast.TryExcept):
                bodies = child.body
                for handler in child.handlers:
                    bodies.extend(handler.body)
                for grandchild in bodies:
                    if isinstance(grandchild, ast.Import):
                        found_try_except_import = True
                    if isinstance(grandchild, ast.Assign):
                        for target in grandchild.targets:
                            if target.id.lower().startswith('has_'):
                                found_has = True
            if found_try_except_import and not found_has:
                self.warnings.append('Found Try/Except block without HAS_ '
                                     'assginment') 

def find_globals(g, tree):
    """Uses AST to find globals in an ast tree"""
    for child in tree:
        if hasattr(child, 'body') and isinstance(child.body, list):
            find_globals(g, child.body)
        elif isinstance(child, (ast.FunctionDef, ast.ClassDef)):
            g.add(child.name)
            continue
        elif isinstance(child, ast.Assign):
            try:
                g.add(child.targets[0].id)
            except (IndexError, AttributeError):
                pass
        elif isinstance(child, ast.Import):
            g.add(child.names[0].name)
        elif isinstance(child, ast.ImportFrom):
            for name in child.names:
                g_name = name.asname or name.name
                if g_name == '*':
                    continue
                g.add(g_name) 

def get_all_namedtuple_names(tree):
    nametuples_names = []
    for node in ast.walk(tree):
        if not isinstance(node, ast.Assign):
            continue
        if not isinstance(node.value, ast.Call):
            continue
        if hasattr(node.value.func, 'id') and node.value.func.id == 'namedtuple':
            nametuples_names.append(node.targets[0].id)
    return nametuples_names 

def generic_visit(self, node):
        if isinstance(node, ast.Call):
            self.functions.append(node)
            self.function_map.update(ast_unparse(node, initial_mapping=True))
        if isinstance(node, ast.Name):
            self.name_types.append(node)
        if isinstance(node, ast.Str):
            self.strings.append(node)
        if isinstance(node, ast.Assign):
            self.assignments.update(ast_unparse(node, initial_mapping=True))
        super(Analyzer, self).generic_visit(node) 

def check_for_wrong_tuple(tree, code, noqa):
    errors = []
    candidates = []
    for assign in ast.walk(tree):
        if not isinstance(assign, (ast.Assign, ast.Return)):
            continue
        elif assign.lineno in noqa:
            continue
        elif isinstance(assign.value, ast.Call):
            continue
        for tuple_el in ast.walk(assign):
            if isinstance(tuple_el, ast.Tuple) and len(tuple_el.elts) == 1:
                candidates.append((assign.lineno, assign.col_offset))
                break
    if not candidates:
        return []
    for candidate in candidates:
        number_nl = 0  # account for logical newlines within statements
        tokens = tokenize.generate_tokens(
            lambda L=iter(code): next(L)
        )
        previous_token = None
        for t in tokens:
            if previous_token is not None and previous_token.type == tokenize.NEWLINE:
                number_nl = 0
            x = TokenInfo(*t)
            if x.start[0] - number_nl != candidate[0]:
                previous_token = x
                continue
            if x.type == tokenize.NL:
                number_nl += 1
            if x.type == token.NEWLINE and ending_of_bad_tuple(previous_token):
                errors.append(x.start)
            if x.type == token.OP and x.string == '=' and previous_token.type != token.NAME:
                x = TokenInfo(*next(tokens))
                if x.type != token.OP and x.string != '(':
                    x_next = TokenInfo(*next(tokens))
                    if ending_of_bad_tuple(x_next):
                        errors.append(x.start)
            previous_token = x
    return errors 

def visit_For(self, loop_):
        """
        >>> self = FirstPassFor(buffer='foo')
        >>> code = '''
        ...
        ... for i in range(5):
        ...     for j in range(5):
        ...         k = i + j
        ...         print(k)
        ...
        ... '''
        >>> tree = ast.parse(code)
        >>> loop_, = tree.body

        """
        loop = self.generic_visit(loop_)
        var = ast.Name(id=__random_string__(), ctx=ast.Store())
        assign = ast.Assign(targets=[var], value=ast.Call(func=ast.Name(id='iter', ctx=ast.Load()), args=[loop.iter], keywords=[]))
        first_pass = ast.Try(
            body=[ast.Assign(targets=[loop.target], value=ast.Call(func=ast.Name(id='next', ctx=ast.Load()), args=[ast.Name(id=var, ctx=ast.Load())], keywords=[]))],
            handlers=[ast.ExceptHandler(type=ast.Name(id='StopIteration', ctx=ast.Load()), name=None, body=[ast.Pass()])],
            orelse=loop.body,
            finalbody=[]
        )
        content = f'`for {astor.to_source(loop.target).strip()} in {astor.to_source(loop.iter).strip()} ...`'
        return [
            make_annotation(buffer=self.buffer, content=content, cell_type='2', lineno=loop.lineno),
            ast.Expr(loop.iter),
            assign,
            first_pass
        ] 

def xproto_check_synchronizer(m):
    try:
        sync_step_path = "synchronizer/steps/sync_%s.py" % m["name"].lower()
        sync_step = open(sync_step_path).read()
    except IOError:
        return "510 Model needs a sync step %s" % sync_step_path

    try:
        sync_step_ast = ast.parse(sync_step)
    except SyntaxError:
        return "511 Could not parse sync step %s" % sync_step_path

    classes = [x for x in sync_step_ast.body if isinstance(x, ast.ClassDef)]
    found_sync_step_class = False

    for c in classes:
        base_names = [v.id for v in c.bases]
        if "SyncStep" in base_names or "SyncInstanceUsingAnsible" in base_names:
            attributes = [x for x in c.body if isinstance(x, ast.Assign)]
            for a in attributes:
                target_names = [t.id for t in a.targets]
                values = a.value.elts if isinstance(a.value, ast.List) else [a.value]
                value_names = [v.id for v in values]

                if "observes" in target_names and m["name"] in value_names:
                    found_sync_step_class = True
                    break

    if not found_sync_step_class:
        return (
            "512 Synchronizer needs a sync step class with an observes field containing %s"
            % m["name"]
        )
    else:
        return "200 OK" 

def visit_Assign(self, assign):
        value_explanation, value_result = self.visit(assign.value)
        explanation = "... = %s" % (value_explanation,)
        name = ast.Name("__exprinfo_expr", ast.Load(),
                        lineno=assign.value.lineno,
                        col_offset=assign.value.col_offset)
        new_assign = ast.Assign(assign.targets, name, lineno=assign.lineno,
                                col_offset=assign.col_offset)
        mod = ast.Module([new_assign])
        co = self._compile(mod, "exec")
        try:
            self.frame.exec_(co, __exprinfo_expr=value_result)
        except Exception:
            raise Failure(explanation)
        return explanation, value_result 

def visit_BoolOp(self, boolop: ast.BoolOp) -> Tuple[ast.Name, str]:
        res_var = self.variable()
        expl_list = self.assign(ast.List([], ast.Load()))
        app = ast.Attribute(expl_list, "append", ast.Load())
        is_or = int(isinstance(boolop.op, ast.Or))
        body = save = self.statements
        fail_save = self.expl_stmts
        levels = len(boolop.values) - 1
        self.push_format_context()
        # Process each operand, short-circuiting if needed.
        for i, v in enumerate(boolop.values):
            if i:
                fail_inner = []  # type: List[ast.stmt]
                # cond is set in a prior loop iteration below
                self.expl_stmts.append(ast.If(cond, fail_inner, []))  # noqa
                self.expl_stmts = fail_inner
            self.push_format_context()
            res, expl = self.visit(v)
            body.append(ast.Assign([ast.Name(res_var, ast.Store())], res))
            expl_format = self.pop_format_context(ast.Str(expl))
            call = ast.Call(app, [expl_format], [])
            self.expl_stmts.append(ast.Expr(call))
            if i < levels:
                cond = res  # type: ast.expr
                if is_or:
                    cond = ast.UnaryOp(ast.Not(), cond)
                inner = []  # type: List[ast.stmt]
                self.statements.append(ast.If(cond, inner, []))
                self.statements = body = inner
        self.statements = save
        self.expl_stmts = fail_save
        expl_template = self.helper("_format_boolop", expl_list, ast.Num(is_or))
        expl = self.pop_format_context(expl_template)
        return ast.Name(res_var, ast.Load()), self.explanation_param(expl) 

def visit_Compare(self, comp: ast.Compare) -> Tuple[ast.expr, str]:
        self.push_format_context()
        left_res, left_expl = self.visit(comp.left)
        if isinstance(comp.left, (ast.Compare, ast.BoolOp)):
            left_expl = "({})".format(left_expl)
        res_variables = [self.variable() for i in range(len(comp.ops))]
        load_names = [ast.Name(v, ast.Load()) for v in res_variables]
        store_names = [ast.Name(v, ast.Store()) for v in res_variables]
        it = zip(range(len(comp.ops)), comp.ops, comp.comparators)
        expls = []
        syms = []
        results = [left_res]
        for i, op, next_operand in it:
            next_res, next_expl = self.visit(next_operand)
            if isinstance(next_operand, (ast.Compare, ast.BoolOp)):
                next_expl = "({})".format(next_expl)
            results.append(next_res)
            sym = BINOP_MAP[op.__class__]
            syms.append(ast.Str(sym))
            expl = "{} {} {}".format(left_expl, sym, next_expl)
            expls.append(ast.Str(expl))
            res_expr = ast.Compare(left_res, [op], [next_res])
            self.statements.append(ast.Assign([store_names[i]], res_expr))
            left_res, left_expl = next_res, next_expl
        # Use pytest.assertion.util._reprcompare if that's available.
        expl_call = self.helper(
            "_call_reprcompare",
            ast.Tuple(syms, ast.Load()),
            ast.Tuple(load_names, ast.Load()),
            ast.Tuple(expls, ast.Load()),
            ast.Tuple(results, ast.Load()),
        )
        if len(comp.ops) > 1:
            res = ast.BoolOp(ast.And(), load_names)  # type: ast.expr
        else:
            res = load_names[0]
        return res, self.explanation_param(self.pop_format_context(expl_call)) 

def test_place_initial_imports(self):
        s = """'Doc string'\nother = stuff"""
        m = rewrite(s)
        assert isinstance(m.body[0], ast.Expr)
        for imp in m.body[1:3]:
            assert isinstance(imp, ast.Import)
            assert imp.lineno == 2
            assert imp.col_offset == 0
        assert isinstance(m.body[3], ast.Assign)
        s = """from __future__ import division\nother_stuff"""
        m = rewrite(s)
        assert isinstance(m.body[0], ast.ImportFrom)
        for imp in m.body[1:3]:
            assert isinstance(imp, ast.Import)
            assert imp.lineno == 2
            assert imp.col_offset == 0
        assert isinstance(m.body[3], ast.Expr)
        s = """'doc string'\nfrom __future__ import division"""
        m = rewrite(s)
        assert isinstance(m.body[0], ast.Expr)
        assert isinstance(m.body[1], ast.ImportFrom)
        for imp in m.body[2:4]:
            assert isinstance(imp, ast.Import)
            assert imp.lineno == 2
            assert imp.col_offset == 0
        s = """'doc string'\nfrom __future__ import division\nother"""
        m = rewrite(s)
        assert isinstance(m.body[0], ast.Expr)
        assert isinstance(m.body[1], ast.ImportFrom)
        for imp in m.body[2:4]:
            assert isinstance(imp, ast.Import)
            assert imp.lineno == 3
            assert imp.col_offset == 0
        assert isinstance(m.body[4], ast.Expr)
        s = """from . import relative\nother_stuff"""
        m = rewrite(s)
        for imp in m.body[:2]:
            assert isinstance(imp, ast.Import)
            assert imp.lineno == 1
            assert imp.col_offset == 0
        assert isinstance(m.body[3], ast.Expr) 

def extract_jsonish_assignments(mod_ast):
  """This extracts all single assignments where the target is a name, and the
  value is a simple 'jsonish' statement (aka Python literal).

  The result is returned as a dictionary of name to the decoded literal.

  Example:
    Foo = "hello"
    Bar = [1, 2, "something"]
    Other, Things = range(2)  # not single assignment
    Bogus = object()  # not a Python literal
    # returns: {"Foo": "hello", "Bar": [1, 2, "something"]}
  """
  ret = {}
  for node in mod_ast.body:
    if not isinstance(node, ast.Assign):
      continue
    if len(node.targets) != 1:
      continue
    if not isinstance(node.targets[0], ast.Name):
      continue
    try:
      ret[node.targets[0].id] = ast.literal_eval(node.value)
    except (KeyError, ValueError):
      pass
  return ret 

def assign(self, expr: ast.expr) -> ast.Name:
        """Give *expr* a name."""
        name = self.variable()
        self.statements.append(ast.Assign([ast.Name(name, ast.Store())], expr))
        return ast.Name(name, ast.Load())