pythonabstract-syntax-treepython-decorators

Is there a way to find undecorated classes/functions in a Python module?


I'm trying to write a tool which will help me find implementation bugs in a large existing python code base that makes heavy use of decorators.

Lets say I have a file, rules.py:

manager = RuleManager()

@manager.register
class A:
    pass

class B:
    pass

Lets say I already know ahead of time what the Manager classes are that use the decorator.

How would I find that class B is undecorated?

Ideally the output would include something like /path/to/rules.py:`B` is undecorated by any manager


Solution

  • Use ast.parse to parse the syntax tree in each file, then fetch all ClassDef's with an empty decorator list and report their relative location:

    from ast import parse, walk, ClassDef, Attribute
    
    # define predicate to test decorator list for presence of @manager.register
    def has_decorator(deco_list):
      for deco in deco_list:
        if isinstance(deco, Attribute) and d.value.id == 'manager' and d.attr == 'register':
          return True
      return False
    
    # parse input text
    with open("/path/to/file.py") as f:
      source_tree = parse(f.read())
    
    # traverse the tree, inspect all class defs for the decorator
    for node in walk(source_tree):
      if isinstance(node, ClassDef) and not has_decorator(node.decorator_list):
        print(f"Class '{node.name}' on line {node.lineno} has no expected decorators")