Printing Python Imports

In my last post, Tips For Pinning Python Requirements Files, I touched briefly on import discovery. I used grep as a crude mechanism to grab python import because it’s Good Enough™. Bash scripts are a great way to get started, but sometimes one requires a more robust tool. I’ve since spent some time diving deeper into the Python import system and discovered a better solution. Let’s start with a single file named myscript.py. We’ll give it a couple of different import flavors to run the gambit.

import requests
from requests import RequestException, HttpError
from requests.api import sessions

r = requests.get("https://duckduckgo.com")
print(r.status_code)

To understand the structure of this script, we can use Python’s ast module The ast module helps Python applications to process trees of the Python abstract syntax grammar. We can build a separate script (called importprinter.py) that will print all the imports it finds:

#!/usr/bin/env python3

import ast
import inspect
import importlib
import sys


class MyNodeVisitor(ast.NodeVisitor):
    """https://docs.python.org/3/library/ast.html#ast.NodeVisitor"""

    def __init__(self):
        self.imports = set()

    def visit_Import(self, node):
        for module_node in node.names:
            self.imports.add(f"import {module_node.name}")
        self.generic_visit(node)

    def visit_ImportFrom(self, node):
        module_name = node.module
        for from_node in node.names:
            self.imports.add(f"from {module_name} import {from_node.name}")
        self.generic_visit(node)

# Consider all arguments passed to be filepaths
for file_path in sys.argv[1:]:
    with open(file_path, "r") as f:
        contents = f.read()
        mod_ast = ast.parse(contents)
        visitor = MyNodeVisitor()
        visitor.visit(mod_ast)
        for import_line in visitor.imports:
            print(import_line)

NodeVisitor is part of the ast module. Each time it visits a certain node it will call the function import_<Class>. We’re looking for both Import and ImportFrom. visit_Import will grab lines that look like import requests. visit_ImportFrom will grab lines that look like from requests import RequestException. Once we’ve grabbed all the imports, we print them. Make sure you run chmod +x importprinter.py before running the following:

$ ./importprinter.py myscript.py
requests.HttpError
requests.api.sessions
requests.RequestException
requests

But, what about multiple python files? We could build logic to traverse directories and grab files in python, but I’ll leave that to bash. By accepting a list of paths via sys.argv[1:], we can can use find and xargs do the lifting.

$ find ~/path/to/pythonfiles -name '*.py' -print0 | xargs -0 ./importprinter.py | sort | uniq

This approach works out cleaner than the previous post’s grep example. Each import is printed on a separate line so we don’t have to worry about brackets. The python script does one thing, print imports of a single file. I prefer using bash tools to handle finding/filtering files. One more handy tool in your toolkit.