Include namespaces (scottrogowski#54)

* --include-only-namespace/function arguments

* fix bugs and add tests

Author: scottrogowski

code2flow/engine.py

@@ -430,6 +430,7 @@ def _find_links(node_a, all_nodes):
 
 
 def map_it(sources, extension, no_trimming, exclude_namespaces, exclude_functions,
+           include_only_namespaces, include_only_functions,
            skip_parse_errors, lang_params):
     '''
     Given a language implementation and a list of filenames, do these things:
@@ -447,6 +448,8 @@ def map_it(sources, extension, no_trimming, exclude_namespaces, exclude_function
     :param bool no_trimming:
     :param list exclude_namespaces:
     :param list exclude_functions:
+    :param list include_only_namespaces:
+    :param list include_only_functions:
     :param bool skip_parse_errors:
     :param LanguageParams lang_params:
 
@@ -475,11 +478,11 @@ def map_it(sources, extension, no_trimming, exclude_namespaces, exclude_function
         file_group = make_file_group(file_ast_tree, source, extension)
         file_groups.append(file_group)
 
-    # 3. Trim namespaces / functions that we don't want
-    if exclude_namespaces:
-        file_groups = _exclude_namespaces(file_groups, exclude_namespaces)
-    if exclude_functions:
-        file_groups = _exclude_functions(file_groups, exclude_functions)
+    # 3. Trim namespaces / functions to exactly what we want
+    if exclude_namespaces or include_only_namespaces:
+        file_groups = _limit_namespaces(file_groups, exclude_namespaces, include_only_namespaces)
+    if exclude_functions or include_only_functions:
+        file_groups = _limit_functions(file_groups, exclude_functions, include_only_functions)
 
     # 4. Consolidate structures
     all_subgroups = flatten(g.all_groups() for g in file_groups)
@@ -557,48 +560,71 @@ def map_it(sources, extension, no_trimming, exclude_namespaces, exclude_function
     return file_groups, all_nodes, edges
 
 
-def _exclude_namespaces(file_groups, exclude_namespaces):
+def _limit_namespaces(file_groups, exclude_namespaces, include_only_namespaces):
     """
     Exclude namespaces (classes/modules) which match any of the exclude_namespaces
 
     :param list[Group] file_groups:
     :param list exclude_namespaces:
+    :param list include_only_namespaces:
     :rtype: list[Group]
     """
+
+    removed_namespaces = set()
+
+    for group in list(file_groups):
+        if group.token in exclude_namespaces:
+            for node in group.all_nodes():
+                node.remove_from_parent()
+            removed_namespaces.add(group.token)
+        if include_only_namespaces and group.token not in include_only_namespaces:
+            for node in group.nodes:
+                node.remove_from_parent()
+            removed_namespaces.add(group.token)
+
+        for subgroup in group.all_groups():
+            print(subgroup, subgroup.all_parents())
+            if subgroup.token in exclude_namespaces:
+                for node in subgroup.all_nodes():
+                    node.remove_from_parent()
+                removed_namespaces.add(subgroup.token)
+            if include_only_namespaces and \
+               subgroup.token not in include_only_namespaces and \
+               all(p.token not in include_only_namespaces for p in subgroup.all_parents()):
+                for node in subgroup.nodes:
+                    node.remove_from_parent()
+                removed_namespaces.add(group.token)
+
     for namespace in exclude_namespaces:
-        found = False
-        for group in list(file_groups):
-            if group.token == namespace:
-                file_groups.remove(group)
-                found = True
-            for subgroup in group.all_groups():
-                if subgroup.token == namespace:
-                    subgroup.remove_from_parent()
-                    found = True
-        if not found:
+        if namespace not in removed_namespaces:
             logging.warning(f"Could not exclude namespace '{namespace}' "
-                            "because it was not found.")
+                             "because it was not found.")
     return file_groups
 
 
-def _exclude_functions(file_groups, exclude_functions):
+def _limit_functions(file_groups, exclude_functions, include_only_functions):
     """
     Exclude nodes (functions) which match any of the exclude_functions
 
     :param list[Group] file_groups:
     :param list exclude_functions:
+    :param list include_only_functions:
     :rtype: list[Group]
     """
+
+    removed_functions = set()
+
+    for group in list(file_groups):
+        for node in group.all_nodes():
+            if node.token in exclude_functions or \
+               (include_only_functions and node.token not in include_only_functions):
+                node.remove_from_parent()
+                removed_functions.add(node.token)
+
     for function_name in exclude_functions:
-        found = False
-        for group in list(file_groups):
-            for node in group.all_nodes():
-                if node.token == function_name:
-                    node.remove_from_parent()
-                    found = True
-        if not found:
+        if function_name not in removed_functions:
             logging.warning(f"Could not exclude function '{function_name}' "
-                            "because it was not found.")
+                             "because it was not found.")
     return file_groups
 
 
@@ -636,6 +662,7 @@ def _generate_final_img(output_file, extension, final_img_filename, num_edges):
 
 def code2flow(raw_source_paths, output_file, language=None, hide_legend=True,
               exclude_namespaces=None, exclude_functions=None,
+              include_only_namespaces=None, include_only_functions=None,
               no_grouping=False, no_trimming=False, skip_parse_errors=False,
               lang_params=None, subset_params=None, level=logging.INFO):
     """
@@ -648,6 +675,8 @@ def code2flow(raw_source_paths, output_file, language=None, hide_legend=True,
     :param bool hide_legend: Omit the legend from the output
     :param list exclude_namespaces: List of namespaces to exclude
     :param list exclude_functions: List of functions to exclude
+    :param list include_only_namespaces: List of namespaces to include
+    :param list include_only_functions: List of functions to include
     :param bool no_grouping: Don't group functions into namespaces in the final output
     :param bool no_trimming: Don't trim orphaned functions / namespaces
     :param bool skip_parse_errors: If a language parser fails to parse a file, skip it
@@ -660,11 +689,16 @@ def code2flow(raw_source_paths, output_file, language=None, hide_legend=True,
 
     if not isinstance(raw_source_paths, list):
         raw_source_paths = [raw_source_paths]
+    lang_params = lang_params or LanguageParams()
+
     exclude_namespaces = exclude_namespaces or []
     assert isinstance(exclude_namespaces, list)
     exclude_functions = exclude_functions or []
     assert isinstance(exclude_functions, list)
-    lang_params = lang_params or LanguageParams()
+    include_only_namespaces = include_only_namespaces or []
+    assert isinstance(include_only_namespaces, list)
+    include_only_functions = include_only_functions or []
+    assert isinstance(include_only_functions, list)
 
     logging.basicConfig(format="Code2Flow: %(message)s", level=level)
 
@@ -691,6 +725,7 @@ def code2flow(raw_source_paths, output_file, language=None, hide_legend=True,
 
     file_groups, all_nodes, edges = map_it(sources, language, no_trimming,
                                            exclude_namespaces, exclude_functions,
+                                           include_only_namespaces, include_only_functions,
                                            skip_parse_errors, lang_params)
 
     if subset_params:
@@ -761,6 +796,12 @@ def main(sys_argv=None):
     parser.add_argument(
         '--exclude-namespaces',
         help='exclude namespaces (Classes, modules, etc) from the output. Comma delimited.')
+    parser.add_argument(
+        '--include-only-functions',
+        help='include only functions in the output. Comma delimited.')
+    parser.add_argument(
+        '--include-only-namespaces',
+        help='include only namespaces (Classes, modules, etc) in the output. Comma delimited.')
     parser.add_argument(
         '--no-grouping', action='store_true',
         help='instead of grouping functions into namespaces, let functions float.')
@@ -801,6 +842,9 @@ def main(sys_argv=None):
 
     exclude_namespaces = list(filter(None, (args.exclude_namespaces or "").split(',')))
     exclude_functions = list(filter(None, (args.exclude_functions or "").split(',')))
+    include_only_namespaces = list(filter(None, (args.include_only_namespaces or "").split(',')))
+    include_only_functions = list(filter(None, (args.include_only_functions or "").split(',')))
+
     lang_params = LanguageParams(args.source_type, args.ruby_version)
     subset_params = SubsetParams.generate(args.target_function, args.upstream_depth,
                                           args.downstream_depth)
@@ -812,6 +856,8 @@ def main(sys_argv=None):
         hide_legend=args.hide_legend,
         exclude_namespaces=exclude_namespaces,
         exclude_functions=exclude_functions,
+        include_only_namespaces=include_only_namespaces,
+        include_only_functions=include_only_functions,
         no_grouping=args.no_grouping,
         no_trimming=args.no_trimming,
         skip_parse_errors=args.skip_parse_errors,

code2flow/model.py

@@ -606,6 +606,15 @@ def remove_from_parent(self):
         if self.parent:
             self.parent.subgroups = [g for g in self.parent.subgroups if g != self]
 
+    def all_parents(self):
+        """
+        Recursively get the entire inheritance tree of this group
+        :rtype: list[Group]
+        """
+        if self.parent:
+            return [self.parent] + self.parent.all_parents()
+        return []
+
     def to_dot(self):
         """
         Returns string format for embedding in a dotfile. Example output:

tests/testdata.py

@@ -81,6 +81,121 @@
             ],
             "expected_nodes": ["simple_b::a", "simple_b::b"]
         },
+
+        {
+            "test_name": "simple_b --include-only-functions",
+            "comment": "Include one a and b",
+            "directory": "simple_b",
+            "kwargs": {"include_only_functions": ["a", "b"]},
+            "expected_edges": [
+                ["simple_b::a", "simple_b::b"],
+                ["simple_b::b", "simple_b::a"]
+            ],
+            "expected_nodes": ["simple_b::a", "simple_b::b"]
+        },
+        {
+            "test_name": "pytz --include-only-functions",
+            "comment": "Include only two functions",
+            "directory": "pytz",
+            "kwargs": {"include_only_functions": ["_fill", "open_resource"]},
+            "expected_edges": [
+                ["__init__::_CountryTimezoneDict._fill", "__init__::open_resource"],
+                ["__init__::_CountryNameDict._fill", "__init__::open_resource"]
+            ],
+            "expected_nodes": ["__init__::_CountryTimezoneDict._fill", "__init__::_CountryNameDict._fill", "__init__::open_resource"]
+        },
+
+        {
+            "test_name": "--include-only-namespaces=reference",
+            "directory": "pytz",
+            "kwargs":{"include_only_namespaces": ["reference"]},
+            "expected_edges": [["reference::USTimeZone.utcoffset",
+                                "reference::USTimeZone.dst"],
+                               ["reference::LocalTimezone.utcoffset",
+                                "reference::LocalTimezone._isdst"],
+                               ["reference::(global)", "reference::USTimeZone.__init__"],
+                               ["reference::USTimeZone.tzname",
+                                "reference::USTimeZone.dst"],
+                               ["reference::LocalTimezone.dst",
+                                "reference::LocalTimezone._isdst"],
+                               ["reference::LocalTimezone.tzname",
+                                "reference::LocalTimezone._isdst"],
+                               ["reference::USTimeZone.dst",
+                                "reference::first_sunday_on_or_after"]],
+            "expected_nodes": ["reference::USTimeZone.utcoffset",
+                               "reference::LocalTimezone.dst",
+                               "reference::LocalTimezone.tzname",
+                               "reference::USTimeZone.__init__",
+                               "reference::USTimeZone.tzname",
+                               "reference::USTimeZone.dst",
+                               "reference::(global)",
+                               "reference::LocalTimezone._isdst",
+                               "reference::LocalTimezone.utcoffset",
+                               "reference::first_sunday_on_or_after"]
+        },
+
+
+        {
+            "test_name": "--include-only-namespaces=USTimeZone",
+            "directory": "pytz",
+            "kwargs":{"include_only_namespaces": ["USTimeZone"]},
+            "expected_edges": [["reference::USTimeZone.utcoffset",
+                                "reference::USTimeZone.dst"],
+                               ["reference::USTimeZone.tzname",
+                                "reference::USTimeZone.dst"]],
+            "expected_nodes": ["reference::USTimeZone.dst",
+                               "reference::USTimeZone.utcoffset",
+                               "reference::USTimeZone.tzname"]
+        },
+        {
+            "test_name": "include_exclude_namespaces",
+            "comment": "Complex including/excluding namespaces",
+            "directory": "pytz",
+            "kwargs": {"include_only_namespaces": ["tzfile","tzinfo"],
+                       "exclude_namespaces": ["DstTzInfo"]},
+            "expected_edges": [["tzinfo::(global)", "tzinfo::memorized_timedelta"],
+                               ["tzfile::build_tzinfo", "tzfile::_std_string"],
+                               ["tzfile::(global)", "tzfile::build_tzinfo"],
+                               ["tzfile::build_tzinfo", "tzinfo::memorized_ttinfo"],
+                               ["tzfile::build_tzinfo", "tzinfo::memorized_timedelta"],
+                               ["tzfile::(global)", "tzfile::_byte_string"],
+                               ["tzinfo::unpickler", "tzinfo::memorized_timedelta"],
+                               ["tzfile::build_tzinfo", "tzfile::_byte_string"],
+                               ["tzfile::build_tzinfo", "tzinfo::memorized_datetime"],
+                               ["tzinfo::memorized_ttinfo",
+                                "tzinfo::memorized_timedelta"]],
+            "expected_nodes": ["tzinfo::unpickler",
+                               "tzfile::build_tzinfo",
+                               "tzinfo::(global)",
+                               "tzfile::_std_string",
+                               "tzinfo::memorized_timedelta",
+                               "tzfile::(global)",
+                               "tzinfo::memorized_datetime",
+                               "tzfile::_byte_string",
+                               "tzinfo::memorized_ttinfo"]
+        },
+        {
+            "test_name": "include_exclude_namespaces_exclude_functions",
+            "comment": "Very complex including/excluding namespaces and excluding functions",
+            "directory": "pytz",
+            "kwargs": {"include_only_namespaces": ["tzfile","tzinfo"],
+                       "exclude_namespaces": ["DstTzInfo"],
+                       "exclude_functions": ["(global)"]},
+            "expected_edges": [["tzfile::build_tzinfo", "tzinfo::memorized_ttinfo"],
+                               ["tzfile::build_tzinfo", "tzinfo::memorized_timedelta"],
+                               ["tzfile::build_tzinfo", "tzinfo::memorized_datetime"],
+                               ["tzinfo::memorized_ttinfo", "tzinfo::memorized_timedelta"],
+                               ["tzfile::build_tzinfo", "tzfile::_std_string"],
+                               ["tzfile::build_tzinfo", "tzfile::_byte_string"],
+                               ["tzinfo::unpickler", "tzinfo::memorized_timedelta"]],
+            "expected_nodes": ["tzinfo::unpickler",
+                               "tzinfo::memorized_datetime",
+                               "tzinfo::memorized_ttinfo",
+                               "tzfile::build_tzinfo",
+                               "tzfile::_std_string",
+                               "tzinfo::memorized_timedelta",
+                               "tzfile::_byte_string"]
+        },
         {
             "test_name": "simple_b --exclude-namespaces not found",
             "comment": "Exclude something not there",