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.
Top comments (0)