How to textually find an imported name in a module

Question

I wrote a method called buildRegex that, given a name (of type str), returns a regex object that finds a from ... import ... nam

Answer 1

I'd recommend that instead of writing complicated regular expressions to parse imports, one would actually use the ast.parse to parse the source code into abstract syntax tree, and find the names from there. asast.parse is guaranteed to parse Python correctly. Something like:

import ast

class ImportFinder(ast.NodeVisitor):
    def __init__(self):
        self.imports = []

    def visit_Import(self, node):
        names = []
        for i in node.names:
            names.append((i.name, i.asname))
        self.imports.append(['import', names])

    def visit_ImportFrom(self, node):
        module = node.module
        level = node.level  # how many dots
        names = []
        for i in node.names:
            names.append((i.name, i.asname))

        self.imports.append(('from', level, module, names))

def parse_imports(source):
    tree = ast.parse(source)
    finder = ImportFinder()
    finder.visit(tree)
    return finder.imports

Example usage:

import pprint

pprint.pprint(parse_imports('''
from foo import bar, baz, frob
from .. import bar as spam, baz as ham, frob
import bar.baz
import bar.foo as baf
'''))

Prints out:

[('from', 0, 'foo', [('bar', None), ('baz', None), ('frob', None)]),
 ('from', 2, None, [('bar', 'spam'), ('baz', 'ham'), ('frob', None)]),
 ['import', [('bar.baz', None)]],
 ['import', [('bar.foo', 'baf')]]]

The integer on the from lines gives the number of . before the module name.

Answer 2

import inspect
import importlib
import ast


class Imports(ast.NodeVisitor):
    def visit_Import(self, node):
        print("In Import")
        for imp in node.names:
            if imp.asname is not None:
                print("module name = {}, alias = {}".format(imp.name, imp.asname))
            else:
                print("module name = {}".format(imp.name))
        print()

    def visit_ImportFrom(self, node):
        print("In ImportFrom")
        for imp in node.names:
            if imp.asname is not None:
                print("module = {}\nname = {}\nalias = {}\nlevel = {}\n".
                      format(node.module, imp.name, imp.asname, node.level))
            else:
                print("module = {}\nname = {}\nlevel = {}\n".
                      format(node.module, imp.name, node.level))
        print()

mod = "temp_test"
mod = importlib.import_module(mod)
p = ast.parse(inspect.getsource(mod))
Imports().visit(p)

Input:

from bisect import bisect_left as bs
import datetime
import time
import numpy as np

def foo():
    from re import findall

class Foo():
    def test(self):
        from re import compile as cp, finditer as ft

Output:

In ImportFrom
module = bisect
name = bisect_left
alias = bs
level = 0


In Import
module name = datetime

In Import
module name = time

In Import
module name = numpy, alias = np

In ImportFrom
module = re
name = findall
level = 0


In ImportFrom
module = re
name = compile
alias = cp
level = 0

module = re
name = finditer
alias = ft
level = 0

class Import(names)

An import statement. names is a list of alias nodes.

class ImportFrom(module, names, level)

Represents from x import y. module is a raw string of the ‘from’ name, without any leading dots, or None for statements such as from . import foo. level is an integer holding the level of the relative import (0 means absolute import).

The greentreesnakes documentation for me at least has a much better explanation of what all the nodes do and how to use the ast module than the actual ast documentation itself.

You can use also pass the module directly or open the py file and pass the content to ast.parse:

with open("temp_test.py") as f:
    p = ast.parse(f.read(), filename="<ast>", mode="exec")
Imports().visit(p)

And passing the module:

import  temp_test

p = ast.parse(inspect.getsource(temp_test))
Imports().visit(p)