Exclusively build for pwdata and filter out classes (#3)

* Add `exclusives` and `skip_classes` to filter out irrelevant modules

pycg/__main__.py

@@ -12,6 +12,10 @@ def main():
     parser.add_argument(
         "--package", help="Package containing the code to be analyzed", default=None
     )
+    parser.add_argument(
+        "--exclusives", help="Exclusive package to analyze", default=[])
+    parser.add_argument(
+        "--skip-classes", help="Classes to skip during traversal", default=[])
     parser.add_argument(
         "--fasten",
         help="Produce call graph using the FASTEN format",
@@ -57,7 +61,9 @@ def main():
     args = parser.parse_args()
 
     cg = CallGraphGenerator(
-        args.entry_point, args.package, args.max_iter, args.operation
+        args.entry_point, args.package,
+        args.exclusives, args.skip_classes,
+        args.max_iter, args.operation
     )
     cg.analyze()
 

pycg/processing/postprocessor.py

@@ -30,6 +30,8 @@ def __init__(
         self,
         input_file,
         modname,
+        exclusives,
+        skip_classes,
         import_manager,
         scope_manager,
         def_manager,
@@ -38,6 +40,8 @@ def __init__(
         modules_analyzed=None,
     ):
         super().__init__(input_file, modname, modules_analyzed)
+        self.exclusives = exclusives
+        self.skip_classes = skip_classes
         self.import_manager = import_manager
         self.scope_manager = scope_manager
         self.def_manager = def_manager
@@ -174,6 +178,10 @@ def visit_FunctionDef(self, node):
         super().visit_FunctionDef(node)
 
     def visit_ClassDef(self, node):
+        # Return if visiting class should be skipped
+        if node.name in self.skip_classes:
+            return
+
         # create a definition for the class (node.name)
         cls_def = self.def_manager.handle_class_def(self.current_ns, node.name)
 
@@ -329,6 +337,8 @@ def update_parent_classes(self, defi):
     def analyze_submodules(self):
         super().analyze_submodules(
             PostProcessor,
+            self.exclusives,
+            self.skip_classes,
             self.import_manager,
             self.scope_manager,
             self.def_manager,

pycg/processing/preprocessor.py

@@ -19,6 +19,7 @@
 # under the License.
 #
 import ast
+import re
 
 from pycg import utils
 from pycg.machinery.definitions import Definition
@@ -30,6 +31,8 @@ def __init__(
         self,
         filename,
         modname,
+        exclusives,
+        skip_classes,
         import_manager,
         scope_manager,
         def_manager,
@@ -40,6 +43,8 @@ def __init__(
         super().__init__(filename, modname, modules_analyzed)
 
         self.modname = modname
+        self.exclusives = exclusives
+        self.skip_classes = skip_classes
         self.mod_dir = "/".join(self.filename.split("/")[:-1])
 
         self.import_manager = import_manager
@@ -48,6 +53,8 @@ def __init__(
         self.class_manager = class_manager
         self.module_manager = module_manager
 
+        self.skip_classes_pattern = re.compile('|'.join(map(re.escape, self.skip_classes)))
+
     def _get_fun_defaults(self, node):
         defaults = {}
         start = len(node.args.args) - len(node.args.defaults)
@@ -71,6 +78,8 @@ def analyze_submodule(self, modname):
         super().analyze_submodule(
             PreProcessor,
             modname,
+            self.exclusives,
+            self.skip_classes,
             self.import_manager,
             self.scope_manager,
             self.def_manager,
@@ -81,6 +90,11 @@ def analyze_submodule(self, modname):
 
     def visit_Module(self, node):
         def iterate_mod_items(items, const):
+            items = [
+                item for item in items 
+                if not self.skip_classes_pattern.search(item)
+            ]
+
             for item in items:
                 defi = self.def_manager.get(item)
                 if not defi:
@@ -203,6 +217,15 @@ def add_external_def(name, target):
         for import_item in node.names:
             src_name = handle_src_name(import_item.name)
             tgt_name = import_item.asname if import_item.asname else import_item.name
+
+            # Limit to exclusive module if exclusives exist
+            if self.exclusives and src_name.split(".")[0] not in self.exclusives:
+                continue
+
+            # Skip specified classes
+            if tgt_name in self.skip_classes:
+                continue
+
             imported_name = self.import_manager.handle_import(src_name, level)
 
             if not imported_name:
@@ -223,6 +246,9 @@ def add_external_def(name, target):
 
         # handle all modules that were not analyzed
         for modname in self.import_manager.get_imports(self.modname):
+            if self.exclusives and modname.split(".")[0] not in self.exclusives:
+                continue
+
             fname = self.import_manager.get_filepath(modname)
 
             if not fname:
@@ -402,6 +428,9 @@ def visit_Lambda(self, node):
 
     def visit_ClassDef(self, node):
         # create a definition for the class (node.name)
+        if node.name in self.skip_classes:
+            return
+
         cls_def = self.def_manager.handle_class_def(self.current_ns, node.name)
 
         mod = self.module_manager.get(self.modname)

pycg/pycg.py

@@ -35,9 +35,13 @@
 
 
 class CallGraphGenerator(object):
-    def __init__(self, entry_points, package, max_iter, operation):
+    def __init__(
+        self, entry_points, package, exclusives, skip_classes, max_iter, operation
+    ):
         self.entry_points = entry_points
         self.package = package
+        self.exclusives = exclusives
+        self.skip_classes = skip_classes
         self.state = None
         self.max_iter = max_iter
         self.operation = operation
@@ -162,6 +166,8 @@ def analyze(self):
         self.do_pass(
             PreProcessor,
             True,
+            self.exclusives,
+            self.skip_classes,
             self.import_manager,
             self.scope_manager,
             self.def_manager,
@@ -179,6 +185,8 @@ def analyze(self):
             self.do_pass(
                 PostProcessor,
                 False,
+                self.exclusives,
+                self.skip_classes,
                 self.import_manager,
                 self.scope_manager,
                 self.def_manager,