def test_handle_import(self):
# test builtin modules
fpath = "input_file.py"
im = ImportManager()
im.set_pkg("")
im.create_node("mod1")
im.set_current_mod("mod1", fpath)
self.assertEqual(im.handle_import("sys", 0), None)
self.assertEqual(im.handle_import("sys", 10), None)
self.assertEqual(im.get_imports("mod1"), set(["sys"]))
# test parent package
class MockImport:
def __init__(self, name):
self.__file__ = name
with mock.patch("importlib.import_module",
return_value=MockImport(
os.path.abspath("mod2.py"))) as mock_import:
modname = im.handle_import("mod2", 0)
self.assertEqual(modname, "mod2")
mock_import.assert_called_with("mod2", package="")
with mock.patch("importlib.import_module",
return_value=MockImport(
os.path.abspath("mod2.py"))) as mock_import:
im.set_current_mod("mod1.mod3", fpath)
modname = im.handle_import("mod2", 1)
self.assertEqual(modname, "mod2")
mock_import.assert_called_once_with(".mod2", package="mod1")
def test_create_edge(self):
fpath = "input_file.py"
im = ImportManager()
im.set_pkg("")
node1 = "node1"
node2 = "node2"
im.create_node(node1)
im.create_node(node2)
# current module not set
with self.assertRaises(ImportManagerError):
im.create_edge(node2)
im.set_current_mod(node1, fpath)
im.create_edge(node2)
self.assertEqual(im.get_imports(node1), set([node2]))
# only non empty strings allowed
with self.assertRaises(ImportManagerError):
im.create_edge("")
with self.assertRaises(ImportManagerError):
im.create_edge(1)
def test_handle_import_level(self):
fpath = "input_file.py"
im = ImportManager()
im.set_pkg("")
im.set_current_mod("mod1.mod2.mod3", fpath)
# gets outside of package scope
with self.assertRaises(ImportError):
im._handle_import_level("something", 4)
self.assertEqual(im._handle_import_level("smth", 2),
("..smth", "mod1"))
self.assertEqual(im._handle_import_level("smth", 1),
(".smth", "mod1.mod2"))
self.assertEqual(im._handle_import_level("smth", 0), ("smth", ""))
def test_hooks(self):
input_file = "somedir/somedir/input_file.py"
im = ImportManager()
im.set_pkg("somedir/somedir")
old_sys_path = copy.deepcopy(sys.path)
old_path_hooks = copy.deepcopy(sys.path_hooks)
custom_loader = "custom_loader"
with mock.patch("importlib.machinery.FileFinder.path_hook",
return_value=custom_loader):
im.install_hooks()
self.assertEqual(sys.path_hooks[0], custom_loader)
self.assertEqual(sys.path[0],
os.path.abspath(os.path.dirname(input_file)))
im.remove_hooks()
self.assertEqual(old_sys_path, sys.path)
self.assertEqual(old_path_hooks, sys.path_hooks)
def test_create_node(self):
fpath = "input_file.py"
im = ImportManager()
im.set_pkg("")
name = "node1"
im.create_node(name)
self.assertEqual(im.get_filepath(name), "")
self.assertEqual(im.get_imports(name), set())
# if a node already exists it can't be added again
with self.assertRaises(ImportManagerError):
im.create_node(name)
# no empty node names
with self.assertRaises(ImportManagerError):
im.create_node("")
with self.assertRaises(ImportManagerError):
im.create_node(1)
def test_custom_loader(self):
fpath = "input_file.py"
im = ImportManager()
im.set_pkg("")
old_sys_path = copy.deepcopy(sys.path)
im.set_current_mod("node1", fpath)
im.create_node("node1")
# an import happens and the loader is called
loader = get_custom_loader(im)("node2", "filepath")
# verify that edges and nodes have been added
self.assertEqual(im.get_imports("node1"), set(["node2"]))
self.assertEqual(im.get_filepath("node2"), os.path.abspath("filepath"))
loader = get_custom_loader(im)("node2", "filepath")
self.assertEqual(im.get_imports("node1"), set(["node2"]))
self.assertEqual(im.get_filepath("node2"), os.path.abspath("filepath"))
self.assertEqual(loader.get_filename("filepath"), "filepath")
self.assertEqual(loader.get_data("filepath"), "")
def test_set_filepath(self):
fpath = "input_file.py"
im = ImportManager()
im.set_pkg("")
name = "node1"
im.create_node(name)
filepath1 = "filepath1"
im.set_filepath(name, filepath1)
self.assertEqual(im.get_filepath(name), os.path.abspath(filepath1))
filepath2 = "filepath2"
im.set_filepath(name, filepath2)
self.assertEqual(im.get_filepath(name), os.path.abspath(filepath2))
# only non empty strings allowed
with self.assertRaises(ImportManagerError):
im.set_filepath(name, "")
with self.assertRaises(ImportManagerError):
im.set_filepath(name, 1)
class CallGraphGenerator(object):
def __init__(self, entry_points, package, max_iter, operation):
self.entry_points = entry_points
self.package = package
self.state = None
self.max_iter = max_iter
self.operation = operation
self.setUp()
def setUp(self):
self.import_manager = ImportManager()
self.scope_manager = ScopeManager()
self.def_manager = DefinitionManager()
self.class_manager = ClassManager()
self.module_manager = ModuleManager()
self.cg = CallGraph()
self.key_errs = KeyErrors()
def extract_state(self):
state = {}
state["defs"] = {}
for key, defi in self.def_manager.get_defs().items():
state["defs"][key] = {
"names": defi.get_name_pointer().get().copy(),
"lit": defi.get_lit_pointer().get().copy()
}
state["scopes"] = {}
for key, scope in self.scope_manager.get_scopes().items():
state["scopes"][key] = set(
[x.get_ns() for (_, x) in scope.get_defs().items()])
state["classes"] = {}
for key, ch in self.class_manager.get_classes().items():
state["classes"][key] = ch.get_mro().copy()
return state
def reset_counters(self):
for key, scope in self.scope_manager.get_scopes().items():
scope.reset_counters()
def has_converged(self):
if not self.state:
return False
curr_state = self.extract_state()
# check defs
for key, defi in curr_state["defs"].items():
if not key in self.state["defs"]:
return False
if defi["names"] != self.state["defs"][key]["names"]:
return False
if defi["lit"] != self.state["defs"][key]["lit"]:
return False
# check scopes
for key, scope in curr_state["scopes"].items():
if not key in self.state["scopes"]:
return False
if scope != self.state["scopes"][key]:
return False
# check classes
for key, ch in curr_state["classes"].items():
if not key in self.state["classes"]:
return False
if ch != self.state["classes"][key]:
return False
return True
def remove_import_hooks(self):
self.import_manager.remove_hooks()
def tearDown(self):
self.remove_import_hooks()
def _get_mod_name(self, entry, pkg):
# We do this because we want __init__ modules to
# only contain the parent module
# since pycg can't differentiate between functions
# coming from __init__ files.
input_mod = utils.to_mod_name(os.path.relpath(entry, pkg))
if input_mod.endswith("__init__"):
input_mod = ".".join(input_mod.split(".")[:-1])
return input_mod
def do_pass(self, cls, install_hooks=False, *args, **kwargs):
modules_analyzed = set()
for entry_point in self.entry_points:
input_pkg = self.package
input_mod = self._get_mod_name(entry_point, input_pkg)
input_file = os.path.abspath(entry_point)
if not input_mod:
continue
if not input_pkg:
input_pkg = os.path.dirname(input_file)
if not input_mod in modules_analyzed:
if install_hooks:
self.import_manager.set_pkg(input_pkg)
self.import_manager.install_hooks()
processor = cls(input_file,
input_mod,
modules_analyzed=modules_analyzed,
*args,
**kwargs)
processor.analyze()
modules_analyzed = modules_analyzed.union(
processor.get_modules_analyzed())
if install_hooks:
self.remove_import_hooks()
def analyze(self):
self.do_pass(PreProcessor, True, self.import_manager,
self.scope_manager, self.def_manager, self.class_manager,
self.module_manager)
self.def_manager.complete_definitions()
iter_cnt = 0
while (self.max_iter < 0
or iter_cnt < self.max_iter) and (not self.has_converged()):
self.state = self.extract_state()
self.reset_counters()
self.do_pass(PostProcessor, False, self.import_manager,
self.scope_manager, self.def_manager,
self.class_manager, self.module_manager)
self.def_manager.complete_definitions()
iter_cnt += 1
self.reset_counters()
if self.operation == utils.constants.CALL_GRAPH_OP:
self.do_pass(CallGraphProcessor,
False,
self.import_manager,
self.scope_manager,
self.def_manager,
self.class_manager,
self.module_manager,
call_graph=self.cg)
elif self.operation == utils.constants.KEY_ERR_OP:
self.do_pass(KeyErrProcessor, False, self.import_manager,
self.scope_manager, self.def_manager,
self.class_manager, self.key_errs)
else:
raise Exception("Invalid operation: " + self.operation)
def output(self):
return self.cg.get()
def output_key_errs(self):
return self.key_errs.get()
def output_edges(self):
return self.key_errors
def output_edges(self):
return self.cg.get_edges()
def _generate_mods(self, mods):
res = {}
for mod, node in mods.items():
res[mod] = {
"filename": os.path.relpath(node.get_filename(), self.package)\
if node.get_filename() else None,
"methods": node.get_methods()
}
return res
def output_internal_mods(self):
return self._generate_mods(self.module_manager.get_internal_modules())
def output_external_mods(self):
return self._generate_mods(self.module_manager.get_external_modules())
def output_functions(self):
functions = []
for ns, defi in self.def_manager.get_defs().items():
if defi.is_function_def():
functions.append(ns)
return functions
def output_classes(self):
classes = {}
for cls, node in self.class_manager.get_classes().items():
classes[cls] = {"mro": node.get_mro(), "module": node.get_module()}
return classes
def get_as_graph(self):
return self.def_manager.get_defs().items()