From 19c8fa389f0b956aa88c1e7a1ec5194f96043c51 Mon Sep 17 00:00:00 2001
From: Ryan Freckleton <ryan.freckleton@gmail.com>
Date: Mon, 12 Feb 2024 13:32:55 -0700
Subject: [PATCH] fix(tree-sitter-typescript-tags): correct tree-sitter queries

Add a new test case for typescript.
Fix queries for typescript.
---
 aider/queries/tree-sitter-typescript-tags.scm | 22 +++----
 tests/test_repomap.py                         | 59 +++++++++++++++++++
 2 files changed, 67 insertions(+), 14 deletions(-)

diff --git a/aider/queries/tree-sitter-typescript-tags.scm b/aider/queries/tree-sitter-typescript-tags.scm
index b1ab874dd..2e6815342 100644
--- a/aider/queries/tree-sitter-typescript-tags.scm
+++ b/aider/queries/tree-sitter-typescript-tags.scm
@@ -1,23 +1,17 @@
-(function_signature
+(function_declaration
   name: (identifier) @name.definition.function) @definition.function
 
-(method_signature
+(method_definition
   name: (property_identifier) @name.definition.method) @definition.method
 
-(abstract_method_signature
-  name: (property_identifier) @name.definition.method) @definition.method
-
-(abstract_class_declaration
+(class_declaration
   name: (type_identifier) @name.definition.class) @definition.class
 
-(module
-  name: (identifier) @name.definition.module) @definition.module
-
 (interface_declaration
-  name: (type_identifier) @name.definition.interface) @definition.interface
+  name: (type_identifier) @name.definition.class) @definition.class
 
-(type_annotation
-  (type_identifier) @name.reference.type) @reference.type
+(type_alias_declaration
+  name: (type_identifier) @name.definition.type) @definition.type
 
-(new_expression
-  constructor: (identifier) @name.reference.class) @reference.class
+(enum_declaration
+  name: (identifier) @name.definition.enum) @definition.enum
diff --git a/tests/test_repomap.py b/tests/test_repomap.py
index a5ce91f1e..96441c6e7 100644
--- a/tests/test_repomap.py
+++ b/tests/test_repomap.py
@@ -1,9 +1,13 @@
+from collections import defaultdict
 import os
 import unittest
+from pathlib import Path
+import networkx as nx
 
 from aider.dump import dump  # noqa: F401
 from aider.io import InputOutput
 from aider.repomap import RepoMap
+from aider import models
 from aider.utils import IgnorantTemporaryDirectory
 
 
@@ -150,5 +154,60 @@ print(my_function(3, 4))
             del repo_map
 
 
+class TestRepoMapTypescript(unittest.TestCase):
+    def test_get_repo_map_typescript(self):
+        # Create a temporary directory with a sample TypeScript file
+        test_file_ts = "test_file.ts"
+        file_content_ts = """\
+interface IMyInterface {
+    someMethod(): void;
+}
+
+type ExampleType = {
+    key: string;
+    value: number;
+};
+
+enum Status {
+    New,
+    InProgress,
+    Completed,
+}
+
+export class MyClass {
+    constructor(public value: number) {}
+
+    add(input: number): number {
+        return this.value + input;
+        return this.value + input;
+    }
+}
+
+export function myFunction(input: number): number {
+    return input * 2;
+}
+"""
+
+        with IgnorantTemporaryDirectory() as temp_dir:
+            with open(os.path.join(temp_dir, test_file_ts), "w") as f:
+                f.write(file_content_ts)
+
+            io = InputOutput()
+            repo_map = RepoMap(root=temp_dir, io=io)
+            other_files = [os.path.join(temp_dir, test_file_ts)]
+            result = repo_map.get_repo_map([], other_files)
+
+            # Check if the result contains the expected tags map with TypeScript identifiers
+            self.assertIn("test_file.ts", result)
+            self.assertIn("IMyInterface", result)
+            self.assertIn("ExampleType", result)
+            self.assertIn("Status", result)
+            self.assertIn("MyClass", result)
+            self.assertIn("add", result)
+            self.assertIn("myFunction", result)
+
+            # close the open cache files, so Windows won't error
+            del repo_map
+
 if __name__ == "__main__":
     unittest.main()