Python, ast, and redbaron

I recently had a use case at work where I wanted to check that file paths given in a Python script actually existed. These paths were in various GitHub repositories, so all I had to do was pull out the paths and check if they exist on GitHub.

There were a few catches though.

First, I couldn’t simply get any string out of each Python script - they needed to be strings specficied by a specific function parameter, and match a regex (e.g., start with ‘abc’).

Second, the script paths lack the GitHub repository root name. This name was part of the function name - so I needed to get access to the function that the path was specified within, and then parse the function name to get the repository name.

The obvious solution I thought was the ast library.

ast library

I started by using ast. The ast.NodeVisitor class seemed like it would do the trick.

An example script (“my_script.py”):

def hello(path, stuff=None):
    return path

if __name__ == " __main__":
    print(hello(path="hello/world.py", stuff="hello mars"))


import ast

class CollectStrings(ast.NodeVisitor):
    def visit_Module(self, node):
        self.out = set()
        self.generic_visit(node)
        return list(filter(lambda w: w.startswith("hello") and w.endswith(".py"), self.out))

    def visit_Str(self, node):
        self.out.add(node.s)

file = "my_script.py"
with open(file, "r") as f:
    body = ast.parse(f.read())

coll = CollectStrings()
coll.visit(body)
## ['hello/world.py']

That worked great at fetching paths - only because all the paths I was looking for started with the same text and all have the same file extension.

HOWEVER - I also needed the function name that the path argument was called from. I tried to make this work with ast.NodeVisitor but couldn’t get it to work.

I eventually got frustrated enough and figured there must be some libraries that build on top of ast that make it easier to work with ast’s in Python.

redbaron

Enter redbaron. I found this library pretty quickly upon searching for a library building on top of ast.

Another example script (“their_script.py”):

def hello(path, stuff=None):
    return path

def goodbye(path, stuff=None):
    return path

def world():
    path_str = hello(path="src/world.py", stuff="hello mars")
    other_path_str = goodbye(path="src/world.py", stuff="hello saturn")

    return path_str, other_path_str

if __name__ == " __main__":
    print(world())


import re
from redbaron import RedBaron

file = "their_script.py"
with open(file, "r") as src:
  red = RedBaron(src.read())

red
## 0 def hello(path, stuff=None):
## return path
##     
##     
##     
## 1 def goodbye(path, stuff=None):
## return path
##     
##     
##     
## 2 def world():
## path_str = hello(path="src/world.py", stuff="hello mars")
## other_path_str = goodbye(path="src/world.py", stuff="hello saturn")
##     
## return path_str, other_path_str
##     
##     
##     
## 3 if __name__ == " __main__":
## print(world())
## 

Even just the resulting object you get from parsing something is useful:

And with .help() you get a very detailed map of the structure of the thing you’re trying to navigate (only printing first 20 lines):

red.help()
## 0 -----------------------------------------------------
## DefNode()
## # identifiers: def, def_, defnode, funcdef, funcdef_
## # default test value: name
## async=False
## name='hello'
## return_annotation ->
## None
## decorators ->
## arguments ->
## * DefArgumentNode()
## # identifiers: def_argument, def_argument_, defargument, defargumentnode
## target ->
## NameNode() ...
## annotation ->
## None
## value ->
## None
## * DefArgumentNode()
## # identifiers: def_argument, def_argument_, defargument, defargumentnode
...

Looking at the result from red.help() I can then use .find_all() to find certain nodes in the ast.

nodes = red.find_all("AtomtrailersNode")
nodes = list(filter(lambda w: "hello" in w.dumps(), nodes))
nodes
## [hello(path="src/world.py", stuff="hello mars"), goodbye(path="src/world.py", stuff="hello saturn")]

Then I can write some okay code to extract out the function name, and ugly code to get the string supplied to the path parameter. Then f-string those together to get the path I’m after.

paths = []
for node in nodes:
    fxn_name = node.name.value
    command = re.search("src/.*\\.py", node.dumps()).group()
    paths.append(f"{fxn_name}/{command}")

for path in paths:
    print(path)
## hello/src/world.py
## goodbye/src/world.py

Not super proud of this but gets the job done for my use case - and when you’re not making open source for others, you don’t need to worry about other use cases :)

I’ll definitely try to learn how to properly extract stuff using redbaron - but it got me to answer much faster than the ast library.