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