pythonpython-2.7mox

Python easy way to read all import statements from py module


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.


Solution

  • 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)