Python ast.AugAssign() Examples

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 assign(self, block, stmt):
    #Evaluate the assignment(s)
    value = stmt.value
    value_variable = None
    for i in range(len(stmt.targets)):
      target = stmt.targets[i]
      if isinstance(target, ast.Tuple):
        #Save intermediate results
        results = []
        #Evaluate individual assignments
        for (t, v) in zip(target.elts, stmt.value.elts):
          results.append(self.assign_single(block, v, t, multi=True))
        #Execute final assignments after intermediate calculations
        for result in results:
          block.add(result)
      else:
        result = self.assign_single(block, value, target, src_variable=value_variable)
        if result is not None:
          block.add(result)
          value_variable = result[-1].expr

  #Parses an augmenting assignment (AST AugAssign) 

def onelinerize(original):
    # original :: string
    # :: string
    t = ast.parse(original)
    table = symtable.symtable(original, '<string>', 'exec')

    original = original.strip()

    # If there's only one line anyways, be lazy
    if len(original.splitlines()) == 1 and \
       len(t.body) == 1 and \
       type(t.body[0]) in (ast.Delete, ast.Assign, ast.AugAssign, ast.Print,
                           ast.Raise, ast.Assert, ast.Import, ast.ImportFrom,
                           ast.Exec, ast.Global, ast.Expr, ast.Pass):
        return original

    return get_init_code(t, table) 

def test_augassign(self):
        aug = ast.AugAssign(ast.Name("x", ast.Load()), ast.Add(),
                            ast.Name("y", ast.Load()))
        self.stmt(aug, "must have Store context")
        aug = ast.AugAssign(ast.Name("x", ast.Store()), ast.Add(),
                            ast.Name("y", ast.Store()))
        self.stmt(aug, "must have Load context") 

def visit_AugAssign(self, node: ast.AugAssign) -> None:
        """
        Visits augmented assignes.

        Raises:
            DoubleMinusOpeationViolation

        """
        self._check_negation(node.op, node.value)
        self._check_string_concat(node.value, node.op)
        self.generic_visit(node) 

def _is_assigned_target(self, node: ast.Subscript) -> bool:
        parent = nodes.get_parent(node)
        if not isinstance(parent, (*AssignNodes, ast.AugAssign)):
            return False
        return any(node == target for target in get_assign_targets(parent)) 

def _check_implicit_sum(self, node: AnyFor) -> None:
        is_implicit_sum = (
            len(node.body) == 1 and
            isinstance(node.body[0], ast.AugAssign) and
            isinstance(node.body[0].op, ast.Add) and
            isinstance(node.body[0].target, ast.Name)
        )
        if is_implicit_sum:
            self.add_violation(ImplicitSumViolation(node)) 

def _check_internals(self, body: Sequence[ast.stmt]) -> None:
        after_closing_node = False
        for index, statement in enumerate(body):
            if after_closing_node:
                self.add_violation(UnreachableCodeViolation(statement))

            if isinstance(statement, self._closing_nodes):
                after_closing_node = True
            elif isinstance(statement, ast.Expr):
                self._check_expression(statement, is_first=index == 0)
            elif isinstance(statement, ast.AugAssign):
                self._check_self_misrefactored_assignment(statement) 

def _check_self_misrefactored_assignment(
        self,
        node: ast.AugAssign,
    ) -> None:
        node_value: ast.expr
        if isinstance(node.value, ast.BinOp):
            node_value = node.value.left

        if isinstance(node.value, self._blocked_self_assignment):
            if name_nodes.is_same_variable(node.target, node_value):
                self.add_violation(MisrefactoredAssignmentViolation(node)) 

def get_assign_targets(
    node: Union[AnyAssignWithWalrus, ast.AugAssign],
) -> List[ast.expr]:
    """Returns list of assign targets without knowing the type of assign."""
    if isinstance(node, (ast.AnnAssign, ast.AugAssign, NamedExpr)):
        return [node.target]
    return node.targets 

def _scope_helper(node):
  """Get the closure of nodes that could begin a scope at this point.

  For instance, when encountering a `(` when parsing a BinOp node, this could
  indicate that the BinOp itself is parenthesized OR that the BinOp's left node
  could be parenthesized.

  E.g.: (a + b * c)   or   (a + b) * c   or   (a) + b * c
        ^                  ^                  ^

  Arguments:
    node: (ast.AST) Node encountered when opening a scope.

  Returns:
    A closure of nodes which that scope might apply to.
  """
  if isinstance(node, ast.Attribute):
    return (node,) + _scope_helper(node.value)
  if isinstance(node, ast.Subscript):
    return (node,) + _scope_helper(node.value)
  if isinstance(node, ast.Assign):
    return (node,) + _scope_helper(node.targets[0])
  if isinstance(node, ast.AugAssign):
    return (node,) + _scope_helper(node.target)
  if isinstance(node, ast.Expr):
    return (node,) + _scope_helper(node.value)
  if isinstance(node, ast.Compare):
    return (node,) + _scope_helper(node.left)
  if isinstance(node, ast.BoolOp):
    return (node,) + _scope_helper(node.values[0])
  if isinstance(node, ast.BinOp):
    return (node,) + _scope_helper(node.left)
  if isinstance(node, ast.Tuple) and node.elts:
    return (node,) + _scope_helper(node.elts[0])
  if isinstance(node, ast.Call):
    return (node,) + _scope_helper(node.func)
  if isinstance(node, ast.GeneratorExp):
    return (node,) + _scope_helper(node.elt)
  if isinstance(node, ast.IfExp):
    return (node,) + _scope_helper(node.body)
  return (node,) 

def visit_AugAssign(self, aug_assign: ast.AugAssign) -> None:
        self.visit_simple_stmt(aug_assign)
        # The variable assigned to is also live (`x` in `x += 5`).
        self._live_vars |= find_variables_by_usage(aug_assign.target)[ast.Store] 

def visit_AugAssign(self, aug_assign: ast.AugAssign) -> ActionsT:
        value, value_actions = self.visit_expr(aug_assign.value)
        target, target_actions = self.visit_expr(aug_assign.target)
        result_node = ast.AugAssign(target=target, op=aug_assign.op, value=value)
        return value_actions + target_actions + [result_node] 

def visit_AugAssign(self, aug_assign: ast.AugAssign, _ctx: CPSTransformerContext) -> VisitReturnT:
        return aug_assign, [] 

def __init__(self, name, source, scope):
        if '__all__' in scope and isinstance(source, ast.AugAssign):
            self.names = list(scope['__all__'].names)
        else:
            self.names = []
        if isinstance(source.value, (ast.List, ast.Tuple)):
            for node in source.value.elts:
                if isinstance(node, ast.Str):
                    self.names.append(node.s)
        super(ExportBinding, self).__init__(name, source) 

def aug_assign(d: ast.AugAssign):
    operator = find(d.op)
    value = find(d.value)
    target = find(d.target)
    return f"{target} = {target} {operator} {value};" 

def gen_statement(self, statement):
        """ Generate code for a statement """
        if isinstance(statement, list):
            for inner_statement in statement:
                self.gen_statement(inner_statement)
        else:
            with self.use_location(statement):
                if isinstance(statement, ast.Pass):
                    pass  # No comments :)
                elif isinstance(statement, ast.Return):
                    self.gen_return(statement)
                elif isinstance(statement, ast.If):
                    self.gen_if(statement)
                elif isinstance(statement, ast.While):
                    self.gen_while(statement)
                elif isinstance(statement, ast.Break):
                    self.gen_break(statement)
                elif isinstance(statement, ast.Continue):
                    self.gen_continue(statement)
                elif isinstance(statement, ast.For):
                    self.gen_for(statement)
                elif isinstance(statement, ast.Assign):
                    self.gen_assign(statement)
                elif isinstance(statement, ast.Expr):
                    self.gen_expr(statement.value)
                elif isinstance(statement, ast.AugAssign):
                    self.gen_aug_assign(statement)
                else:  # pragma: no cover
                    self.not_impl(statement) 

def visit_AugAssign(self, node: ast.AugAssign) -> None:
        # +=, -=, /=, *=
        print(f"AugAssign: {node}")
        print(ast.dump(node))
        self.generic_visit(node) 

def p_expr_stmt_2(p):
    '''expr_stmt : testlist augassign testlist'''
    #                     1         2        3
    ctx_to_store(p[1])
    p[0] = ast.AugAssign(p[1], p[2], p[3], rule=inspect.currentframe().f_code.co_name)
    inherit_lineno(p[0], p[1]) 

def p_expr_stmt_1(p):
    '''expr_stmt : testlist augassign yield_expr'''
    #                     1         2          3
    ctx_to_store(p[1])
    p[0] = ast.AugAssign(p[1], p[2], p[3], rule=inspect.currentframe().f_code.co_name)
    inherit_lineno(p[0], p[1]) 

def statement(self, block, stmt):
    #Add a comment
    self.add_comment(block, stmt)
    #Parse the statement
    try:
      if isinstance(stmt, ast.Assign):
        self.assign(block, stmt)
      elif isinstance(stmt, ast.Return):
        self.return_(block, stmt)
      elif isinstance(stmt, ast.Expr):
        self.docstring_(block, stmt)
      elif isinstance(stmt, ast.If):
        self.if_(block, stmt)
      elif isinstance(stmt, ast.While):
        self.while_(block, stmt)
      elif isinstance(stmt, ast.AugAssign):
        self.augassign(block, stmt)
      else:
        Parser._dump(stmt, 'Unexpected Statement')
        raise Exception('Unexpected Statement (%s)'%(stmt.__class__))
    except:
      line = stmt.lineno
      src = self.source.split('\n')[line - 1].strip()
      print('Line %d: %s'%(line, src))
      raise

  #===========================================================
  # Public interface
  #===========================================================
  #Parses the kernel using the specified live function 

def __init__(self, name, source, scope):
        if '__all__' in scope and isinstance(source, ast.AugAssign):
            self.names = list(scope['__all__'].names)
        else:
            self.names = []
        if isinstance(source.value, (ast.List, ast.Tuple)):
            for node in source.value.elts:
                if isinstance(node, ast.Str):
                    self.names.append(node.s)
        super(ExportBinding, self).__init__(name, source) 

def __init__(self, name, source, scope):
        if '__all__' in scope and isinstance(source, ast.AugAssign):
            self.names = list(scope['__all__'].names)
        else:
            self.names = []

        def _add_to_names(container):
            for node in container.elts:
                if isinstance(node, ast.Str):
                    self.names.append(node.s)

        if isinstance(source.value, (ast.List, ast.Tuple)):
            _add_to_names(source.value)
        # If concatenating lists
        elif isinstance(source.value, ast.BinOp):
            currentValue = source.value
            while isinstance(currentValue.right, ast.List):
                left = currentValue.left
                right = currentValue.right
                _add_to_names(right)
                # If more lists are being added
                if isinstance(left, ast.BinOp):
                    currentValue = left
                # If just two lists are being added
                elif isinstance(left, ast.List):
                    _add_to_names(left)
                    # All lists accounted for - done
                    break
                # If not list concatenation
                else:
                    break
        super(ExportBinding, self).__init__(name, source) 

def test_augassign(self):
        aug = ast.AugAssign(ast.Name("x", ast.Load()), ast.Add(),
                            ast.Name("y", ast.Load()))
        self.stmt(aug, "must have Store context")
        aug = ast.AugAssign(ast.Name("x", ast.Store()), ast.Add(),
                            ast.Name("y", ast.Store()))
        self.stmt(aug, "must have Load context") 

def test_augassign(self):
        aug = ast.AugAssign(ast.Name("x", ast.Load()), ast.Add(),
                            ast.Name("y", ast.Load()))
        self.stmt(aug, "must have Store context")
        aug = ast.AugAssign(ast.Name("x", ast.Store()), ast.Add(),
                            ast.Name("y", ast.Store()))
        self.stmt(aug, "must have Load context") 

def __init__(self, name, source, scope):
        if '__all__' in scope and isinstance(source, ast.AugAssign):
            self.names = list(scope['__all__'].names)
        else:
            self.names = []
        if isinstance(source.value, (ast.List, ast.Tuple)):
            for node in source.value.elts:
                if isinstance(node, ast.Str):
                    self.names.append(node.s)
        super(ExportBinding, self).__init__(name, source) 

def check_code(self, code, nodes):
        linestarts = dict(dis.findlinestarts(code))
        instructions = get_instructions(code)
        lineno = None
        for inst in instructions:
            if time.time() - self.start_time > 45 * 60:
                # Avoid travis time limit of 50 minutes
                raise TimeOut

            lineno = linestarts.get(inst.offset, lineno)
            if not inst.opname.startswith((
                    'BINARY_', 'UNARY_', 'LOAD_ATTR', 'LOAD_METHOD', 'LOOKUP_METHOD',
                    'SLICE+', 'COMPARE_OP', 'CALL_', 'IS_OP', 'CONTAINS_OP',
            )):
                continue
            frame = C()
            frame.f_lasti = inst.offset
            frame.f_code = code
            frame.f_globals = globals()
            frame.f_lineno = lineno
            source = Source.for_frame(frame)
            node = None

            try:
                try:
                    node = Source.executing(frame).node
                except Exception:
                    if inst.opname.startswith(('COMPARE_OP', 'CALL_')):
                        continue
                    if isinstance(only(source.statements_at_line(lineno)), (ast.AugAssign, ast.Import)):
                        continue
                    raise
            except Exception:
                print(source.text, lineno, inst, node and ast.dump(node), code, file=sys.stderr, sep='\n')
                raise

            nodes[node].append((inst, frame.__dict__))

            yield [inst.opname, node_string(source, node)]

        for const in code.co_consts:
            if isinstance(const, type(code)):
                for x in self.check_code(const, nodes):
                    yield x 

def visit_Call(self, node):
        _id = get_call_names_as_string(node.func)
        local_definitions = self.module_definitions_stack[-1]

        alias = handle_aliases_in_calls(_id, local_definitions.import_alias_mapping)
        if alias:
            definition = local_definitions.get_definition(alias)
        else:
            definition = local_definitions.get_definition(_id)

        # e.g. "request.args.get" -> "get"
        last_attribute = _id.rpartition('.')[-1]

        if definition:
            if definition in self.function_definition_stack:
                log.debug("Recursion encountered in function %s", _id)
                return self.add_blackbox_or_builtin_call(node, blackbox=True)
            if isinstance(definition.node, ast.ClassDef):
                self.add_blackbox_or_builtin_call(node, blackbox=False)
            elif isinstance(definition.node, ast.FunctionDef):
                self.undecided = False
                self.function_return_stack.append(_id)
                self.function_definition_stack.append(definition)
                return self.process_function(node, definition)
            else:
                raise Exception('Definition was neither FunctionDef or ' +
                                'ClassDef, cannot add the function ')
        elif (
            not self._within_mutating_call and
            last_attribute in MUTATORS
            and isinstance(node.func, ast.Attribute)
        ):
            # Change list.append(x) ---> list += list.append(x)
            # This does in fact propagate as we don't know that append returns None
            fake_aug_assign = ast.AugAssign(
                target=node.func.value,
                op=ast.Add,
                value=node,
            )
            ast.copy_location(fake_aug_assign, node)
            self._within_mutating_call = True  # Don't do this recursively
            result = self.visit(fake_aug_assign)
            self._within_mutating_call = False
            return result
        elif last_attribute not in BUILTINS:
            # Mark the call as a blackbox because we don't have the definition
            return self.add_blackbox_or_builtin_call(node, blackbox=True)
        return self.add_blackbox_or_builtin_call(node, blackbox=False) 

def visit_AugAssign(self, node: ast.AugAssign) -> ast.AST:
        """AugAssign is ``-=, +=, /=, *=`` for augmented assignment."""
        self.generic_visit(node)
        log_header = f"visit_AugAssign: {self.src_file}:"

        # custom mapping of string keys to ast operations that can be used
        # in the nodes since these overlap with BinOp types
        aug_mappings = {
            "AugAssign_Add": ast.Add,
            "AugAssign_Sub": ast.Sub,
            "AugAssign_Mult": ast.Mult,
            "AugAssign_Div": ast.Div,
        }

        rev_mappings = {v: k for k, v in aug_mappings.items()}
        idx_op = rev_mappings.get(type(node.op), None)

        # edge case protection in case the mapping isn't known for substitution
        # in that instance, return the node and take no action
        if not idx_op:
            LOGGER.debug(
                "%s (%s, %s): unknown aug_assignment: %s",
                log_header,
                node.lineno,
                node.col_offset,
                type(node.op),
            )
            return node

        node_span = NodeSpan(node)
        idx = LocIndex(
            ast_class="AugAssign",
            lineno=node_span.lineno,
            col_offset=node_span.col_offset,
            op_type=idx_op,
            end_lineno=node_span.end_lineno,
            end_col_offset=node_span.end_col_offset,
        )

        self.locs.add(idx)

        if idx == self.target_idx and self.mutation in aug_mappings and not self.readonly:
            LOGGER.debug("%s mutating idx: %s with %s", log_header, self.target_idx, self.mutation)
            return ast.copy_location(
                ast.AugAssign(
                    target=node.target,
                    op=aug_mappings[self.mutation](),  # awkward syntax to call type from mapping
                    value=node.value,
                ),
                node,
            )

        LOGGER.debug("%s (%s, %s): no mutations applied.", log_header, node.lineno, node.col_offset)
        return node