diff --git a/aider/repomap.py b/aider/repomap.py index c82b944dd..4862670d3 100644 --- a/aider/repomap.py +++ b/aider/repomap.py @@ -8,6 +8,7 @@ from pathlib import Path import networkx as nx import tiktoken from diskcache import Cache +from grep_ast import TreeContext from tqdm import tqdm from tree_sitter_languages import get_language, get_parser @@ -16,7 +17,7 @@ from aider.parsers import filename_to_lang from .dump import dump # noqa: F402 -Tag = namedtuple("Tag", "fname line name kind".split()) +Tag = namedtuple("Tag", "fname rel_fname line name kind".split()) def to_tree(tags): @@ -25,26 +26,41 @@ def to_tree(tags): tags = sorted(tags) + cur_fname = None + context = None output = "" - last = [None] * len(tags[0]) - tab = "\t" for tag in tags: - tag = list(tag) + if type(tag) is tuple: + this_fname = tag[0] + else: + this_fname = tag.rel_fname - for i in range(len(last) + 1): - if i == len(last): - break - if last[i] != tag[i]: - break + if this_fname != cur_fname: + if context: + context.add_context() + output += cur_fname + ":\n" + output += context.format() + context = None + elif cur_fname: + output += cur_fname + "\n" - num_common = i + if type(tag) is not tuple: + context = TreeContext( + tag.rel_fname, + Path(tag.fname).read_text(), # TODO: encoding + color=False, + line_number=False, + child_context=False, + last_line=False, + margin=0, + mark_lois=False, + header_pad=1, + loi_pad=0, + ) + cur_fname = this_fname - indent = tab * num_common - rest = tag[num_common:] - for item in rest: - output += indent + str(item) + "\n" - indent += tab - last = tag + if context: + context.add_lines_of_interest([tag.line]) return output @@ -151,7 +167,7 @@ class RepoMap: except FileNotFoundError: self.io.tool_error(f"File not found error: {fname}") - def get_tags(self, fname): + def get_tags(self, fname, rel_fname): lang = filename_to_lang(fname) if not lang: return @@ -168,8 +184,8 @@ class RepoMap: return query_scm = query_scm.read_text() - code = Path(fname).read_text() - tree = parser.parse(bytes(code, "utf8")) + code = Path(fname).read_text() # TODO: encoding + tree = parser.parse(bytes(code, "utf-8")) # Run the tags queries query = language.query(query_scm) @@ -186,9 +202,10 @@ class RepoMap: continue result = Tag( - name=node.text.decode("utf-8"), # TODO: encoding? - kind=kind, fname=fname, + rel_fname=rel_fname, + name=node.text.decode("utf-8"), + kind=kind, line=node.start_point[0], ) @@ -222,7 +239,7 @@ class RepoMap: personalization[rel_fname] = 1.0 chat_rel_fnames.add(rel_fname) - tags = self.get_tags(fname) + tags = self.get_tags(fname, rel_fname) if tags is None: continue @@ -235,8 +252,9 @@ class RepoMap: if tag.kind == "ref": references[tag.name].append(rel_fname) - dump(definitions) - dump(references) + ## + # dump(definitions) + # dump(references) idents = set(defines.keys()).intersection(set(references.keys()))