fix: forgot to 'poetry add click' as a dependency!

Author: chriscarrollsmith

codegraph/core.py

@@ -1,7 +1,7 @@
 import os
 from argparse import Namespace
-from collections import defaultdict
-from typing import Dict, List, Text, Tuple
+from collections import defaultdict, deque
+from typing import Dict, List, Set, Text, Tuple
 
 from codegraph.parser import Import, create_objects_array
 from codegraph.utils import get_python_paths_list
@@ -36,7 +36,7 @@ def get_code_objects(paths_list: List) -> Dict:
 
 class CodeGraph:
     def __init__(self, args: Namespace):
-        self.paths_list = get_python_paths_list(args.paths)
+        self.paths_list = get_python_paths_list(args["paths"])
         # get py modules list data
         self.modules_data = get_code_objects(self.paths_list)
 
@@ -87,6 +87,42 @@ def usage_graph(self) -> Dict:
         dependencies = populate_free_nodes(self.modules_data, dependencies)
         return dependencies
 
+    def get_dependencies(self, file_path: str, distance: int) -> Dict[str, Set[str]]:
+        """
+        Get dependencies that are 'distance' nodes away from the given file.
+
+        :param file_path: Path of the file to start from
+        :param distance: Number of edges to traverse
+        :return: Dictionary with distances as keys and sets of dependent files as values
+        """
+        dependencies = {i: set() for i in range(1, distance + 1)}
+        graph = self.usage_graph()
+
+        if file_path not in graph:
+            return dependencies
+
+        queue = deque([(file_path, 0)])
+        visited = set()
+
+        while queue:
+            current_file, current_distance = queue.popleft()
+
+            if current_distance >= distance:
+                continue
+
+            if current_file not in visited:
+                visited.add(current_file)
+
+                for entity, used_entities in graph[current_file].items():
+                    for used_entity in used_entities:
+                        if "." in used_entity:
+                            dependent_file = used_entity.split(".")[0] + ".py"
+                            if dependent_file != current_file:
+                                dependencies[current_distance + 1].add(dependent_file)
+                                queue.append((dependent_file, current_distance + 1))
+
+        return dependencies
+
 
 def get_module_name(code_path: Text) -> Text:
     module_name = os.path.basename(code_path).replace(".py", "")

codegraph/main.py

@@ -19,7 +19,9 @@
     is_flag=True,
     help="Don't visualize code dependencies as graph",
 )
-def cli(paths, object_only):
+@click.option("--file-path", help="File path to start dependency search from")
+@click.option("--distance", type=int, help="Distance to search for dependencies")
+def cli(paths, object_only, file_path, distance):
     """
     Tool that creates a graph of code to show dependencies between code entities (methods, classes, etc.).
     CodeGraph does not execute code, it is based only on lex and syntax parsing.
@@ -33,18 +35,30 @@ def cli(paths, object_only):
         )
         sys.exit(1)
 
-    args = {"paths": paths, "object_only": object_only}
+    args = {
+        "paths": paths,
+        "object_only": object_only,
+        "file_path": file_path,
+        "distance": distance,
+    }
     main(args)
 
 
 def main(args):
-    usage_graph = core.CodeGraph(args).usage_graph()
-    pprint.pprint(usage_graph)
-    if not args["object_only"]:
-        # to make more quick work if not needed to visualize
-        import codegraph.vizualyzer as vz
+    code_graph = core.CodeGraph(args)
+    usage_graph = code_graph.usage_graph()
 
-        vz.draw_graph(usage_graph)
+    if args.get("file_path") and args.get("distance"):
+        dependencies = code_graph.get_dependencies(args["file_path"], args["distance"])
+        print(f"Dependencies for {args['file_path']}:")
+        for distance, files in dependencies.items():
+            print(f"  Distance {distance}: {', '.join(files)}")
+    else:
+        pprint.pprint(usage_graph)
+        if not args["object_only"]:
+            import codegraph.vizualyzer as vz
+
+            vz.draw_graph(usage_graph)
 
 
 if __name__ == "__main__":

poetry.lock

@@ -21,6 +21,7 @@ classifiers = [
 python = "^3.12"
 matplotlib = "^3.9"
 networkx = "^3.4"
+click = "^8.1.7"
 
 [tool.poetry.dev-dependencies]
 pytest = "^8.3"