Better than grep

The Visitor

We're interested in one feature of the abstract syntax tree of Python: the ast.Constant objects where the constant's value is a string. Ideally, this will focus on those strings that start with a SQL keyword.

Here's the starting point:

class StringVisitor(ast.NodeVisitor):
    def __init__(self) -> None:
        self.sql_blocks = []
    def visit_Constant(self, node: ast.AST) -> None:
        match node.value:
            case str() if sql_like(node.value):
                self.sql_blocks.append(range(node.lineno+1, node.end_lineno))
            case _:  # Ignore all other types
                pass

The visit_Constant() method is invoked for each ast.Constant object in the tree. This will find nodes that are strings and pass the sql_like() filter.

This net is a little too fine. It can will capture docstring comments that happen to look SQL-like. It's essentially the same as the grep output with two improvements:

• No baffling regular expression.
• The output can be a more useful data structure, not a file of lines of text.

While nicer in some respects, it's incomplete. The next step is to add the required feature to ignore docstring comments. There two choices:

1. Write a better sql_like() function.
2. Exclude any string constant that is the first line of the body of a module, class, or function definition.

Option 1 can can involve finding a SQL parser to see if a string really is pure SQL. Or, it can require writing a better regular expression to locate likely SQL statements.

Here's the unit test case:

>>> src_text = """
... def f(c):
...     '''
...         Select the right answers.
...     '''
...     r = c.execute('''
...         SELECT * FROM DUMMY
...     ''')
... """

The docstring comment starts with a SQL keyword. Ugh. It seems kind of daunting to locate a suitable SQL parser. The regular expression to distinguish casual use of SQL-like keywords seems hopeless complicated. There's something better: exclude docstring constants.

Exclusion Rules

The ast.NodeVisitor implementation has a handy feature. This permits an application to choose to visit or skip the subsidiary nodes of an object. When a class overrides a visit_XXX() method, the override can call the self.generic_visit(node) to visit all the children. If the overriding method does not evaluate self.generic_visit(node), the children are not examined.

This is a bit too strict.

Skipping all children of a module, class definition or function definition isn't helpful. The rest of the children could have SQL code. It's important to skip only the very first line of code when this is a string constant. The rest of the code needs to be visited.

I decided to accumulate an "ignore these" set of nodes. This set of nodes will be the first line of code that's also a string constant. The visit_Constant() can then politely decline these nodes.

Here's the visitor class looks:

class StringVisitor(ast.NodeVisitor):
    def __init__(self) -> None:
        self.docstrings = set()  # Docstring nodes at the start of module, class, def
        self.sql_blocks = []
    def exclude_docstring(self, node: ast.AST) -> None:
        if isinstance(node.body[0], ast.Expr) and isinstance(node.body[0].value, ast.Constant):
            self.docstrings.add(node.body[0].value)
        self.generic_visit(node)
    visit_Module = exclude_docstring
    visit_ClassDef = exclude_docstring
    visit_FunctionDef = exclude_docstring
    def visit_Constant(self, node: ast.AST) -> None:
        match node.value:
        case str() if sql_like(node.value):
            if node not in self.docstrings:
                self.sql_blocks.append(range(node.lineno+1, node.end_lineno))
            case _:  # Ignore all other types
                pass

First, I've added a a set of nodes to exclude. (Nodes are immutable, and have a hash value, that's why a set works well for this.)

Second, there's a generic visit function, exclude_docstring(), that handles ast.Module, ast.ClassDef and ast.FunctionDef classes. I can define the needed visti_XXX() methods for these classes to all use the generic exclusion method.

Finally, I need to make sure any ast.Constant node wasn't already excluded because it was the first string in a definition.

With this, I can now parse the source, and apply the visitor. The sql_blocks attribute will have a list of range() objects that point to the SQL statements.

The Parsing and Visiting

Here's the final bit.

tree = ast.parse(src_text, "<test>")
sv = StringVisitor()
sv.visit(tree)
print(sv.sql_blocks)

This will parse the code, then visit the ast.

When this is done, it has a list of range() objects. These ranges can be used to highlight the proper lines of code from the source file. Your application may be different, of course, you may want to simply write the literals to a TOML configuration file, for example.

"Why," you might ask, "do you have range() objects instead of the code?"