Python easy way to read all import statements from py module

Question

I am trying to create a helper function to read a file and mock out all imports for a unit test. I have to read the file vs import since i dont have those things on python path.

Example code:

#module.py
import com.stackoverflow.question
from com.stackoverflow.util import test_func
from com.stackoverflow.util import TestClass

#magic helper: what i want
magic = process('<path_to>/module.py')
for module in magic.modules_as_strings():
    #todo  would have to recuirsively add each path
    # so i would first create com, then com.stackoverflow, etc
    setattr(self, module, StubModules(module)
for obj in magic.sink:
    #these would be "from"  from x import Y
    #its basically just creating self.Y = object
    setattr(self, object)

Above is the mock code, I am really looking for the best way to just tokenize the file for "from/import statements"

That make sense? I know I could read the file line by line, but I was hoping for a cleaner/concise way.

Let me know if you have any questions.

Could you add an example output of the magic calls (data structures)? — GaretJax, Jan 28 '12 at 23:16
sorry i missed this but i was just expecting the string fully qualified module so something like a.b.c.module_name . I was going to tokenize it.. — Nix, Jan 28 '12 at 23:44
I hope my solution will work as well. You can always strip out the unnecessary pieces. — GaretJax, Jan 28 '12 at 23:46

GaretJax · Accepted Answer · 2012-01-28T23:38:23.373

32

Using the AST module, it is pretty easy:

import ast
from collections import namedtuple

Import = namedtuple("Import", ["module", "name", "alias"])

def get_imports(path):
    with open(path) as fh:        
       root = ast.parse(fh.read(), path)

    for node in ast.iter_child_nodes(root):
        if isinstance(node, ast.Import):
            module = []
        elif isinstance(node, ast.ImportFrom):  
            module = node.module.split('.')
        else:
            continue

        for n in node.names:
            yield Import(module, n.name.split('.'), n.asname)

For a module like this:

from coco import bunny
from coco.bungy import carrot
from meta import teta
from rocket import spaceship as sp
import bingo
import com.stackoverflow
import motorbike as car
import module1, module2

s="a random variable"

def func():
    """And a function"""

The output is:

>>> for imp in get_imports("/path/to/file.py"): print imp
Import(module=['coco'], name=['bunny'], alias=None)
Import(module=['coco', 'bungy'], name=['carrot'], alias=None)
Import(module=['meta'], name=['teta'], alias=None)
Import(module=['rocket'], name=['spaceship'], alias='sp')
Import(module=[], name=['bingo'], alias=None)
Import(module=[], name=['com', 'stackoverflow'], alias=None)
Import(module=[], name=['motorbike'], alias='car')
Import(module=[], name=['module1'], alias=None)
Import(module=[], name=['module2'], alias=None)

edited Jan 28 '12 at 23:38

answered Jan 28 '12 at 23:30

GaretJax

7,462
1
38
47

Hmmph. I beat you by 81 seconds but yours is prettier. +1. – DSM Jan 28 '12 at 23:35
Thanks! ;-) I was waiting for a reply from the OP and eventually gave up when I saw that you've already replied… :-P – GaretJax Jan 28 '12 at 23:37
9

You should use `ast.walk` if you have code that does this sort of thing: if something: import xxxx – Tom Tanner Jan 19 '16 at 11:19
3

Thankfully, `ast.walk` is a straight swap for `ast.iter_child_nodes` in this code. Walk is recursive, iter_child_nodes is not, and will not drop down a single block level (never mind infinitely). So, along with conditional imports, this is also an issue for code e.g. `try: import xxxx`. Walk will report that import, iter_child_nodes will not. – BuvinJ Apr 19 '20 at 19:47
Any way to exclude the python standard library? – j7skov Jun 14 '23 at 16:53

Python easy way to read all import statements from py module

1 Answers1

Linked