Mirror of forum.nim-lang.org

11205 :: AST walking similar to python

• riscv (orginal) [2024-03-12T22:30:52+01:00] view original

  In Python, it is trivial to do this kind of AST walking for any valid python code that is provided as a string in a variable.

  import ast
  
  code_string = """
  def foo(x, y):
      z = x + y
      return f(z, g(z))
  """
  module_node = ast.parse(code_string)
  
  class FunctionVisitor(ast.NodeVisitor):
      def visit_FunctionDef(self, node):
          print(f"Found function declaration: {node.name}")
          for stmt in node.body:
              self.visit(stmt)
      
      def visit_Assign(self, node):
          print(f"Found assignment to: {node.targets[0].id}")
          self.visit(node.value)
      
      def visit_Call(self, node):
          print(f"Found function call: {node.func.id}")
          for arg in node.args:
              self.visit(arg)
      
      def visit_BinOp(self, node):
          print(f"Found binary operation: {node.op.__class__.__name__}")
          self.visit(node.left)
          self.visit(node.right)
  
  visitor = FunctionVisitor()
  visitor.visit(module_node)
  Run

  I saw the macros section and found some relevant functionality such as parseExpr and parseStmt but I am not sure what to use/combine to achieve the same thing. Any help would be appreciated.

---

• riscv (orginal) [2024-03-12T23:37:42+01:00] view original

  Something similar to this partly works:

  import macros
  
  proc parseCustomStmt(code: string): NimNode {.compileTime.} =
    let tree = parseStmt(code)
    
    tree
  
  macro parseAndVisit(code: static[string]): void =
    let tree = parseCustomStmt(code)
    
    proc visit(n: NimNode): void =
      case n.kind
      of nnkStmtList, nnkBlockStmt:
        for child in n:
          visit(child)
      
      of nnkProcDef:
        echo "Found procedure declaration: ", n[0].strVal
        for child in n[2]:
          visit(child)
      
      of nnkAsgn:
        echo "Found assignment to: ", n[0].strVal
        visit(n[1])
      
      of nnkCall:
        echo "Found function call to: ", n[0].strVal
        for arg in n[1..^1]:
          visit(arg)
      
      else:
        discard
    
    visit(tree)
  
  parseAndVisit("""
  proc foo(x, y: int): int =
    let z = x
    return z
  """)
  Run

  The nodes in the function body are not triggered though. Not sure why.

---

• riscv (orginal) [2024-03-13T00:59:53+01:00] view original

  Some nodes were missing, found everything:

  import macros
  
  proc parseCustomStmt(code: string): NimNode {.compileTime.} =
    let tree = parseStmt(code)
    return tree
  
  macro parseAndVisit(code: static[string]): void =
    let tree = parseCustomStmt(code)
    
    proc visit(n: NimNode): void =
      case n.kind
      of nnkStmtList, nnkBlockStmt:
        for child in n:
          visit(child)
      
      of nnkProcDef:
        echo "Found procedure declaration: ", n[0].strVal
        for child in n:
          visit(child)
      
      of nnkVarSection, nnkLetSection:
        for vars in n:
          if vars.kind == nnkIdentDefs:
            for varDef in vars:
              if varDef.kind == nnkIdent:
                echo "Found variable declaration: ", varDef.strVal
                visit(varDef)
      
      of nnkReturnStmt:
        echo "Found return statement"
        if n.len > 0 and n[0].kind != nnkEmpty: visit(n[0])
      
      else:
        discard
    
    visit(tree)
  
  parseAndVisit("""
  proc foo(x, y: int): int =
    let z = x
    return z
  """)
  Run

---