Use AST analysis to find which modules can be lazy imported

lazy_import_analyzer.py

import sys
import ast
import importlib.util

def safe_assign(value):
    safe = {ast.Num, ast.NameConstant, ast.Str}
    return type(value) in safe

class Transformer(ast.NodeTransformer):
    def __init__(self, *args, **kwargs):
        ast.NodeTransformer.__init__(self, *args, **kwargs)
        self.is_lazy = True
        self.force_lazy = False
        self.imports = []

    def visit_Module(self, node):
        for stmt in node.body:
            if isinstance(stmt, ast.ImportFrom):
                self.imports.append(stmt.module)
            elif isinstance(stmt, ast.Import):
                for alias in stmt.names:
                    self.imports.append(alias.name)
            elif isinstance(stmt, ast.Assign):
                targets = stmt.targets
                if len(targets) == 1 and targets[0].id == '__lazy__':
                    self.force_lazy = True
                else:
                    if not safe_assign(stmt.value):
                        print('assign', stmt.value)
                        self.is_lazy = False
            elif isinstance(stmt, ast.FunctionDef):
                pass
            elif isinstance(stmt, ast.ClassDef):
                pass
            else:
                print('non-lazy %r' % stmt)
                self.is_lazy = False
        return self.generic_visit(node)


def parse(buf, filename='<string>'):
    if isinstance(buf, bytes):
        buf = importlib.util.decode_source(buf)
    try:
        node = ast.parse(buf, filename)
    except SyntaxError as e:
        # set the filename attribute
        raise SyntaxError(str(e), (filename, e.lineno, e.offset, e.text))
    t = Transformer()
    t.visit(node)
    return t.is_lazy, t.force_lazy, t.imports


def main():
    for fn in sys.argv[1:]:
        with open(fn) as fp:
            buf = fp.read()
            is_lazy, force, imports = parse(buf)
            print('lazy = %s force_lazy = %s' % (is_lazy, force))
            for name in imports:
                print('import', name)

if __name__ == '__main__':
    main()