Python 3.x: how to use ast to search for a print statement

Question

I'm creating a test which is supposed to check if a function contains a print statement or not (Python 3.x, I'm using 3.7.4). I've been using ast to check for similar things (referencing the answer in this question), such as return or list comprehensions, but I'm getting stuck on print.

An online AST explorer lists a Print subclass in the body, and it's taking Python 3 prints, so I know it's not a Python 2 thing.

The Green Tree Snakes ast docs say that Print only has an ast node in Python 2. This is closer to what I'm experiencing. Here's a function that I was going to use to make an assertion:

def printsSomething(func):
    return any(isinstance(node, ast.Print) for node in ast.walk(ast.parse(inspect.getsource(func))))

returns:

TypeError: isinstance() arg 2 must be a type or tuple of types

I'm assuming this has to do with print being a function in Python 3.x, but I can't figure out how to use this knowledge to my advantage. How would I use ast to find out if print has been called?

I would like to reiterate that I've gotten this code to work for other ast nodes such as return, so I should be confident that it's not a bug specific to my code.

Thanks!

score 4 · Answer 1 · answered Sep 05 '19 at 14:51

print is a function in python 3 so you need to check for an ast.Expr which contains an ast.Call for which the ast.Name has the id print.

Here's a simple function:

def bar(x: str) -> None:
    string = f"Hello {x}!"  # ast.Assign
    print(string)           # ast.Expr

Heres's the full ast dump:

Module(body=[FunctionDef(name='bar', args=arguments(args=[arg(arg='x', annotation=Name(id='str', ctx=Load()))], vararg=None, kwonlyargs=[], kw_defaults=[], kwarg=None, defaults=[]), body=[Assign(targets=[Name(id='string', ctx=Store())], value=JoinedStr(values=[Str(s='Hello '), FormattedValue(value=Name(id='x', ctx=Load()), conversion=-1, format_spec=None), Str(s='!')])), Expr(value=Call(func=Name(id='print', ctx=Load()), args=[Name(id='string', ctx=Load())], keywords=[]))], decorator_list=[], returns=NameConstant(value=None))])

The relevant part (for print) is:

Expr(value=Call(func=Name(id='print', ctx=Load())

Below is a simple example with a node visitor (sublcassing ast.NodeVisitor):

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import ast
import inspect
from typing import Callable


class MyNodeVisitor(ast.NodeVisitor):
    def visit_Expr(self, node: ast.Expr):
        """Called when the visitor visits an ast.Expr"""
        print(f"Found expression node at: line: {node.lineno}; col: {node.col_offset}")

        # check "value" which must be an instance of "Call" for a 'print'
        if not isinstance(node.value, ast.Call):
            return

        # now check the function itself.
        func = node.value.func  # ast.Name
        if func.id == "print":
            print("found a print")


def contains_print(f: Callable):
    source = inspect.getsource(f)
    node = ast.parse(source)
    func_name = [_def.name for _def in node.body if isinstance(_def, ast.FunctionDef)][0]
    print(f"{'-' * 79}\nvisiting function: {func_name}")
    print(f"node dump: {ast.dump(node)}")
    node_visitor = MyNodeVisitor()
    node_visitor.visit(node)


def foo(x: int) -> int:
    return x + 1

def bar(x: str) -> None:
    string = f"Hello {x}!"  # ast.Assign
    print(string)           # ast.Expr

def baz(x: float) -> float:
    if x == 0.0:
        print("oh noes!")
        raise ValueError

    return 10 / x


if __name__ == "__main__":
    contains_print(bar)
    contains_print(foo)
    contains_print(baz)

Here the output (minus the ast dumps):

-------------------------------------------------------------------------------
visiting function: bar
Found expression node at: line: 3; col: 4
found a print
-------------------------------------------------------------------------------
visiting function: foo
-------------------------------------------------------------------------------
visiting function: baz
Found expression node at: line: 3; col: 8
found a print

kd88 · Accepted Answer · 2019-09-06T09:58:11.617

If any named object (including functions such as print) is called then there will be at least one _ast.Name object amongst your nodes. The name of the object ('print') is stored under the id attribute of this node.

As I'm sure you're aware, print changed from a statement to a function between python version 2 and 3, possibly explaining why you were running into problems.

Try the following:

import ast
import inspect

def do_print():
    print('hello')

def dont_print():
    pass

def prints_something(func):
    is_print = False
    for node in ast.walk(ast.parse(inspect.getsource(func))):
        try:
            is_print = (node.id == 'print')
        except AttributeError:  # only expect id to exist for Name objs
            pass
        if is_print:
            break
    return is_print

prints_something(do_print), prints_something(dont_print)

>>> True, False

...or if you're a fan of one-liners (where func is your function to test):

any(hasattr(node,'id') and node.id == 'print' 
    for node in ast.walk(ast.parse(inspect.getsource(func))))

I was really hoping for a one-liner, and have spent the last hour or so trying to refactor, but I can't hack around the need for the try/except. I really like this answer since it works with the style in the question. Thanks! — The Renaissance, Sep 05 '19 at 15:47

Python 3.x: how to use ast to search for a print statement

2 Answers2