Python ast.copy_location() 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 visit_Attribute( self, node ):
    self.generic_visit( node )
    if isinstance( node._object, _SignalSlice ):
      if node._object.slice.step:
        raise VerilogTranslationError(
          'Slices with steps ([start:stop:step]) are not translatable!\n',
          node.lineno
        )
      new_node = ast.Subscript( node.value,
                   ast.Slice( ast.Num( node._object.slice.start ),
                              ast.Num( node._object.slice.stop ),
                              None ),
                   None,
                 )
      new_node._object = node._object
      return ast.copy_location( new_node, node )
    return node

#-------------------------------------------------------------------------
# InferTemporaryTypes
#------------------------------------------------------------------------- 

def visit_Attribute( self, node ):

    # Visit children
    self.generic_visit( node )

    # If the direct child of this attribute was a portbundle then the node
    # will be removed by the visitor. We must update our name to include
    # portbundle name for proper mangling.
    if self.portbundle:
      new_node = _ast.Name( id  = '{}_{}'.format(self.portbundle, node.attr ),
                            ctx = node.ctx )
      new_node._object = node._object
      node = new_node

    # Attribute is a PortBundle, remove the node, set the submodule name
    if isinstance( node._object, PortBundle ):
      self.portbundle = node.attr
      return None

    # Otherwise, clear the submodule name, return node unmodified
    self.portbundle = None
    return ast.copy_location( node, node ) 

def visit_BinOp(self, node: ast.BinOp) -> Union[ast.BinOp, ast.Num]:
        """Evaluate binary operation and return ast.Num if operands are bound.

        Args:
            node: Binary operation to evaluate.

        Returns:
            Evaluated value.
        """
        node.left = self.visit(node.left)
        node.right = self.visit(node.right)
        if isinstance(node.left, ast.Num) and isinstance(node.right, ast.Num):
            val = ast.Num(n=self._match_ops(node.op, self._binary_ops,
                                            node.left.n, node.right.n))
            return ast.copy_location(val, node)
        return node 

def visit_Call(self, node: ast.Call) -> Union[ast.Call, ast.Num]:
        """Evaluate function and return ast.Num if all arguments are bound.

        Args:
            node: Function to evaluate.

        Returns:
            Evaluated value.

        Raises:
            PulseError: When unsupported or unsafe function is specified.
        """
        if not isinstance(node.func, ast.Name):
            raise PulseError('Unsafe expression is detected.')
        node.args = [self.visit(arg) for arg in node.args]
        if all(isinstance(arg, ast.Num) for arg in node.args):
            if node.func.id not in self._math_ops.keys():
                raise PulseError('Function %s is not supported.' % node.func.id)
            _args = [arg.n for arg in node.args]
            _val = self._math_ops[node.func.id](*_args)
            if not _val.imag:
                _val = _val.real
            val = ast.Num(n=_val)
            return ast.copy_location(val, node)
        return node 

def visit_Call(self, node):
        if node.starargs:
            star = gast.Starred(self._visit(node.starargs), gast.Load())
            gast.copy_location(star, node)
            star.end_lineno = star.end_col_offset = None
            starred = [star]
        else:
            starred = []

        if node.kwargs:
            kwargs = [gast.keyword(None, self._visit(node.kwargs))]
        else:
            kwargs = []

        new_node = gast.Call(
            self._visit(node.func),
            self._visit(node.args) + starred,
            self._visit(node.keywords) + kwargs,
        )
        gast.copy_location(new_node, node)
        new_node.end_lineno = new_node.end_col_offset = None
        return new_node 

def visit_Name( self, node ):
    new_node = LocalVar( id=node.id )
    new_node._name = node.id
    return ast.copy_location( new_node, node )


#-------------------------------------------------------------------------
# ReorderAST
#-------------------------------------------------------------------------
# Reorders an AST branch beginning with the indicated node.  Intended
# for inverting the order of Name/Attribute chains so that the Name
# node comes first, followed by chains of Attribute/Subscript nodes.
#
# This visitor will also insert Self nodes to represent references to the
# self variable, and remove Index nodes which seem to serve no useful
# purpose. 

def visit_BoolOp(self, node: ast.BoolOp) -> ast.AST:
        """Boolean operations, AND/OR."""
        self.generic_visit(node)
        log_header = f"visit_BoolOp: {self.src_file}:"

        node_span = NodeSpan(node)
        idx = LocIndex(
            ast_class="BoolOp",
            lineno=node_span.lineno,
            col_offset=node_span.col_offset,
            op_type=type(node.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 and not self.readonly:
            LOGGER.debug("%s mutating idx: %s with %s", log_header, self.target_idx, self.mutation)
            return ast.copy_location(ast.BoolOp(op=self.mutation(), values=node.values), node)

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

def visit_Ellipsis(self, node):
        new_node = gast.Constant(
            Ellipsis,
            None,
        )
        gast.copy_location(new_node, node)
        new_node.end_lineno = new_node.end_col_offset = None
        return new_node 

def visit_Str(self, node):
        new_node = gast.Constant(
            node.s,
            None,
        )
        gast.copy_location(new_node, node)
        new_node.end_lineno = new_node.end_col_offset = None
        return new_node 

def visit_Subscript(self, node):
        new_slice = self._visit(node.slice)
        if isinstance(node.slice, ast.Ellipsis):
            new_slice = gast.Index(new_slice)

        new_node = gast.Subscript(
            self._visit(node.value),
            new_slice,
            self._visit(node.ctx),
        )
        gast.copy_location(new_node, node)
        new_node.end_lineno = new_node.end_col_offset = None
        return new_node 

def visit_ClassDef(self, node):
        new_node = ast.ClassDef(
            self._visit(node.name),
            self._visit(node.bases),
            self._visit(node.body),
            self._visit(node.decorator_list),
        )

        ast.copy_location(new_node, node)
        return new_node 

def visit_Num(self, node):
        new_node = gast.Constant(
            node.n,
            None,
        )
        gast.copy_location(new_node, node)
        new_node.end_lineno = new_node.end_col_offset = None
        return new_node 

def visit_comprehension(self, node):
        new_node = gast.comprehension(
            target=self._visit(node.target),
            iter=self._visit(node.iter),
            ifs=self._visit(node.ifs),
            is_async=0,
        )
        gast.copy_location(new_node, node)
        new_node.end_lineno = new_node.end_col_offset = None
        return new_node

    # arguments 

def visit_Raise(self, node):
        if isinstance(node.exc, gast.Call) and \
           isinstance(node.exc.func, gast.Attribute) and \
           node.exc.func.attr == 'with_traceback':
            raised = self._visit(node.exc.func.value)
            traceback = self._visit(node.exc.args[0])
        else:
            raised = self._visit(node.exc)
            traceback = None
        new_node = ast.Raise(raised, None, traceback)
        ast.copy_location(new_node, node)
        return new_node 

def visit_Attribute(self, node):
        """Flatten one level of attribute access."""
        new_node = ast.Name("%s.%s" % (node.value.id, node.attr), node.ctx)
        return ast.copy_location(new_node, node) 

def visit_Call( self, node ):
    self.generic_visit( node )
    if isinstance(node.func, _ast.Name) and node.func.id == "copy":
      return ast.copy_location( node.args[0], node )
    return node

#-----------------------------------------------------------------------
# InferTypes
#----------------------------------------------------------------------- 

def visit_Attribute( self, node ):
    self.generic_visit( node )
    if isinstance(node.value, _ast.Subscript):
      subscript       = node.value
      node.value      = subscript.value
      subscript.value = node
      return ast.copy_location( subscript, node )

    return node 

def visit_Attribute( self, node ):

    # If a parent node is going to be removed, update this name

    # SubModel List
    if self.attr   and isinstance( node._object[0], Model ):
      node.attr = '{}${}'.format( node.attr, self.attr )
      node._object = PortList([ getattr( x, self.attr ) for x in node._object ])
      node._object.name = node.attr

    # PortBundle List
    elif self.attr and isinstance( node._object[0], PortBundle ):
      node.attr = '{}_{}'.format( node.attr, self.attr )
      node._object = PortList([ getattr( x, self.attr ) for x in node._object ])
      node._object.name = node.attr

    # Unknown!!!
    elif self.attr:
      raise Exception( "Don't know how to flatten this!" )

    # Exit early if theres no value attribute
    if not hasattr( node, 'value' ):
      return node

    # If the child is a subscript node, this node will be removed
    if isinstance( node.value, ast.Subscript ):
      self.attr = node.attr
      self.generic_visit( node )
      self.attr = None
      node = node.value

    return ast.copy_location( node, node ) 

def visit_Attribute( self, node ):
    self.generic_visit( node )
    if isinstance( node._object, slice ):
      if node._object.step:
        raise VerilogTranslationError(
          'Slices with steps ([start:stop:step]) are not translatable!\n',
          node.lineno
        )
      new_node = ast.Slice( ast.Num( node._object.start ),
                            ast.Num( node._object.stop ),
                            None )
      return ast.copy_location( new_node, node )
    return node 

def visit_FunctionDef( self, node ):
    #self.generic_visit( node ) # visit children, uneeded?

    # TODO: currently assume decorator is of the form 'self.dec_name'
    if len( node.decorator_list ) != 1:
      def get_dec_name( dec ):
        if   hasattr( dec, 'id'   ): return dec.id
        elif hasattr( dec, 'attr' ): return dec.attr
        else:                        return dec
      raise VerilogTranslationError(
        'Translated behavioral blocks should only have one decorator!\n'
        'Current decorators: {}'.format(
          [ get_dec_name( x ) for x in node.decorator_list ]
        ), node.lineno
      )

    dec = node.decorator_list[0].attr

    # create a new FunctionDef node that deletes the decorators
    new_node = ast.FunctionDef( name=node.name, args=node.args,
                                body=node.body, decorator_list=[dec])

    return ast.copy_location( new_node, node )

#-------------------------------------------------------------------------
# ThreeExprLoops
#-------------------------------------------------------------------------
# Replace calls to range()/xrange() in a for loop with a Slice object
# describing the bounds (upper/lower/step) of the iteration. 

def visit_Module( self, node ):
    #self.generic_visit( node ) # visit children, uneeded?

    # copy the function body, delete module references
    return ast.copy_location( node.body[0], node )

#-------------------------------------------------------------------------
# SimplifyDecorator
#-------------------------------------------------------------------------
# Make the decorator contain text strings, not AST Trees 

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 visit_Attribute( self, node ):

    if node.attr in ['next', 'value', 'n', 'v']:
      # Update the Load/Store information
      node.value.ctx = node.ctx
      return ast.copy_location( node.value, node )

    return node

#-------------------------------------------------------------------------
# RemoveSelf
#-------------------------------------------------------------------------
# Remove references to self.
# TODO: make Attribute attached to self a Name node? 

def visit_Call( self, node ):
    # func, args, keywords, starargs, kwargs
    # Check that this is just a normal function call, not something weird

    self.generic_visit( node )

    if node.func.id == 'range':
      if   len( node.args ) == 1:
        start = _ast.Num( n=0 )
        stop  = node.args[0]
        step  = _ast.Num( n=1 )
      elif len( node.args ) == 2:
        start = node.args[0]
        stop  = node.args[1]
        step  = _ast.Num( n=1 ) # TODO: should be an expression
      elif len( node.args ) == 3:
        start = node.args[0]
        stop  = node.args[1]
        step  = node.args[2]
      else:
        raise Exception("Invalid # of arguments to range function!")

      new_node = _ast.Slice( lower=start, upper=stop, step=step )

    else:
      new_node = node

    return ast.copy_location( new_node, node )

#------------------------------------------------------------------------
# PyObj
#------------------------------------------------------------------------ 

def visit_Index( self, node ):
    # Remove Index nodes, they seem pointless
    child = self.visit( node.value )
    return ast.copy_location( child, node )

  #-----------------------------------------------------------------------
  # visit_Call
  #-----------------------------------------------------------------------
  # Specially handle certain function calls 

def visit_Subscript( self, node ):

    # Visit the object being sliced
    new_value = self.visit( node.value )

    # Visit the index of the slice; stash and restore the current_obj
    stash = self.current_obj
    self.current_obj = None
    new_slice = self.visit( node.slice )
    self.current_obj = stash

    # If current_obj not initialized, it is a local temp. Don't replace.
    if   not self.current_obj:
      new_node = _ast.Subscript( value=new_value, slice=new_slice, ctx=node.ctx )
    # If current_obj is a Bits object, replace with a BitSlice node.
    elif isinstance( self.current_obj.inst, (Bits, InPort, OutPort) ):
      new_node = BitSlice( value=new_value, slice=new_slice, ctx=node.ctx )
    # If current_obj is a list object, replace with an ArrayIndex node.
    elif isinstance( self.current_obj.inst, list ):
      new_node = ArrayIndex( value=new_value, slice=new_slice, ctx=node.ctx )
      # TODO: Want to do this for lists, but can't add attribute
      #       handling in translation instead
      #self.current_obj.inst.name = self.current_obj.inst[0].name.split('[')[0]
    # Otherwise, throw an exception
    else:
      print( self.current_obj )
      raise Exception("Unknown type being subscripted!")

    # Update the current_obj to contain the obj returned by subscript
    # TODO: check that type of all elements in item are identical
    # TODO: won't work for lists that are initially empty
    # TODO: what about lists that initially contain None?
    if self.current_obj:
      self.current_obj.update( '[]', self.current_obj.inst[0] )
    node._object = self.current_obj.inst if self.current_obj else None

    return ast.copy_location( new_node, node )

  #-----------------------------------------------------------------------
  # visit_Index
  #----------------------------------------------------------------------- 

def visit_Name( self, node ):

    # If the name is not in closed_vars, it is a local temporary
    if   node.id not in self.closed_vars:
      new_node = Temp( id=node.id )
      new_obj  = None

    # If the name points to the model, this is a reference to self (or s)
    elif self.closed_vars[ node.id ] is self.model:
      new_node = Self( id=node.id )
      new_obj  = PyObj( '', self.closed_vars[ node.id ] )

    # Otherwise, we have some other variable captured by the closure...
    # TODO: should we allow this?
    else:
      new_node = node
      new_obj  = PyObj( node.id, self.closed_vars[ node.id ] )

    # Store the new_obj
    self.current_obj = new_obj
    node._object = self.current_obj.inst if self.current_obj else None

    # Return the new_node
    return ast.copy_location( new_node, node )

  #-----------------------------------------------------------------------
  # visit_Subscript
  #----------------------------------------------------------------------- 

def visit_Module( self, node ):
    # visit children
    self.generic_visit( node )

    # copy the function body, delete module references
    return ast.copy_location( node.body[0], node )

  #-----------------------------------------------------------------------
  # visit_FunctionDef
  #----------------------------------------------------------------------- 

def visit_FunctionDef( self, node ):
    # store self_name
    # TODO: fix so that this properly detects the "self" param
    if node.args.args: s = node.args.args[0].id
    else:              s = 's'
    self.self_name   = s
    # visit children
    self.generic_visit( node )
    # simplified decorator list
    decorator_list = [ x._name for x in node.decorator_list
                       if isinstance( x, Self ) ]
    # create a new FunctionDef node
    new_node = ast.FunctionDef( name=node.name, args=node.args,
                                body=node.body,
                                decorator_list=decorator_list
                                #decorator_list=[]
                               )
    # create a new function that deletes the decorators
    return new_node

  #def visit_Attribute( self, node ):
  #  reverse_branch = ReorderAST(self.self_name).reverse( node )
  #  return ast.copy_location( reverse_branch, node )
  #def visit_Subscript( self, node ):
  #  reverse_branch = ReorderAST(self.self_name).reverse( node )
  #  return ast.copy_location( reverse_branch, node )
  #def visit_Index( self, node ):
  #  reverse_branch = ReorderAST(self.self_name).reverse( node )
  #  return ast.copy_location( reverse_branch, node )