refac

aider/repomap.py

@@ -2,13 +2,13 @@ import colorsys
 import json
 import os
 import random
-
-# import shelve
 import subprocess
 import sys
 import tempfile
 from collections import Counter, defaultdict
 
+# import shelve
+import graphviz
 import networkx as nx
 import tiktoken
 from pygments.lexers import guess_lexer_for_filename
@@ -263,48 +263,11 @@ class RepoMap:
         res = [token[2] for token in tokens if token[1] in Token.Name]
         return res
 
-
-def find_py_files(directory):
-    if not os.path.isdir(directory):
-        return [directory]
-
-    py_files = []
-    for root, dirs, files in os.walk(directory):
-        for file in files:
-            if file.endswith(".py"):
-                py_files.append(os.path.join(root, file))
-    return py_files
-
-
-def get_random_color():
-    hue = random.random()
-    r, g, b = [int(x * 255) for x in colorsys.hsv_to_rgb(hue, 1, 0.75)]
-    res = f"#{r:02x}{g:02x}{b:02x}"
-    return res
-
-
-def call_map():
-    import graphviz
-
-    fnames = sys.argv[1:]
-
-    fnames = []
-    for dname in sys.argv[1:]:
-        fnames += find_py_files(dname)
-
-    fnames = sorted(fnames)
-
-    rm = RepoMap()
-
-    # res = rm.get_tags_map(fnames)
-    # print(res)
-
+    def get_ranked_tags(self, fnames):
         defines = defaultdict(set)
         references = defaultdict(list)
         definitions = defaultdict(set)
 
-    root = os.path.commonpath(fnames)
-
         personalization = dict()
 
         show_fnames = set()
@@ -351,10 +314,6 @@ def call_map():
 
         for ident in idents:
             definers = defines[ident]
-        num_defs = len(definers)
-        # if num_defs > 3:
-        #    continue
-
             for referencer, num_refs in Counter(references[ident]).items():
                 for definer in definers:
                     if referencer == definer:
@@ -387,8 +346,6 @@ def call_map():
                 ident = data["ident"]
                 ranked_definitions[(dst, ident)] += data["rank"]
 
-    dot = graphviz.Digraph(graph_attr={"ratio": ".5"})
-
         clusters = dict()
         for fname in set(show_fnames):
             clusters[fname] = graphviz.Digraph(f"cluster_{fname}")
@@ -408,11 +365,10 @@ def call_map():
 
             ranked_tags += list(definitions.get((fname, ident), []))
 
-    N = 100
-    ranked_tags = ranked_tags[:N]
-    tree = to_tree(ranked_tags)
-    print(tree)
-    dump(len(tree))
+        draw_graph = False
+
+        if draw_graph:
+            dot = graphviz.Digraph(graph_attr={"ratio": ".5"})
 
             for cluster in clusters.values():
                 dot.subgraph(cluster)
@@ -428,11 +384,51 @@ def call_map():
                     # penwidth=str(weight), color=color, fontcolor=color, label=label,
                 )
 
-    ###
-    # dot.render("tmp", format="pdf", view=True)
-    return
+            dot.render("tmp", format="pdf", view=True)
+
+        return ranked_tags
+
+    def get_ranked_tags_map(self, fnames):
+        ranked_tags = self.get_ranked_tags(fnames)
+
+        N = 100
+        ranked_tags = ranked_tags[:N]
+        tree = to_tree(ranked_tags)
+
+        return tree
+
+
+def find_py_files(directory):
+    if not os.path.isdir(directory):
+        return [directory]
+
+    py_files = []
+    for root, dirs, files in os.walk(directory):
+        for file in files:
+            if file.endswith(".py"):
+                py_files.append(os.path.join(root, file))
+    return py_files
+
+
+def get_random_color():
+    hue = random.random()
+    r, g, b = [int(x * 255) for x in colorsys.hsv_to_rgb(hue, 1, 0.75)]
+    res = f"#{r:02x}{g:02x}{b:02x}"
+    return res
 
 
 if __name__ == "__main__":
-    call_map()
-    # print(rm.get_tags_map(sys.argv[1:]))
+    fnames = sys.argv[1:]
+
+    fnames = []
+    for dname in sys.argv[1:]:
+        fnames += find_py_files(dname)
+
+    fnames = sorted(fnames)
+
+    root = os.path.commonpath(fnames)
+
+    rm = RepoMap(root=root)
+    repo_map = rm.get_ranked_tags_map(fnames)
+    print(repo_map)
+    print(len(repo_map))