fix: handle decorated identifiers

elifarley · elifarley · commit 7e82824ecab9 · 2024-12-11T02:30:23.000-03:00
diff --git a/src/cedarscript_editor/tree_sitter_identifier_finder.py b/src/cedarscript_editor/tree_sitter_identifier_finder.py
@@ -128,7 +128,8 @@ def find_identifiers(
         match identifier_type:
             case 'method':
                 identifier_type = 'function'
-        candidate_nodes = self.language.query(self.query_info[identifier_type].format(name=name)).captures(self.tree.root_node)
+        _query = self.query_info[identifier_type].format(name=name)
+        candidate_nodes = self.language.query(_query).captures(self.tree.root_node)
         if not candidate_nodes:
             return []
         # Convert captures to boundaries and filter by parent
@@ -198,9 +199,9 @@ def parents(self) -> list[ParentInfo]:
         current = self.node.parent
 
         while current:
-            # Check if current node is a container type we care about
-            if current.type.endswith('_definition'):
-                # Try to find the name node - exact field depends on language
+            # Check if current node is a container type we care about - TODO exact field depends on language
+            if current.type.endswith('_definition') and current.type != 'decorated_definition':
+                # Try to find the name node - TODO exact field depends on language
                 name = None
                 for child in current.children:
                     if child.type == 'identifier' or child.type == 'name':
@@ -242,14 +243,13 @@ def associate_identifier_parts(captures: Iterable[CaptureInfo], lines: Sequence[
                 raise ValueError(f'Parent node not found for [{capture.capture_type} - {capture.node_type}] ({capture.node.text.decode("utf-8").strip()})')
             match capture_type:
                 case 'body':
-                    parent = parent._replace(body=range_spec)
+                    parent.body=range_spec
                 case 'docstring':
-                    parent = parent._replace(docstring=range_spec)
+                    parent.docstring=range_spec
                 case 'decorator':
-                    parent = parent.decorators.append(range_spec)
+                    parent.append_decorator(range_spec)
                 case _ as invalid:
                     raise ValueError(f'Invalid capture type: {invalid}')
-            identifier_map[parent_key] = parent
 
     return sorted(identifier_map.values(), key=lambda x: x.whole.start)
 
@@ -260,6 +260,8 @@ def find_parent_definition(node):
     while node.parent:
         node = node.parent
         if node.type.endswith('_definition'):
+            if node.type == 'decorated_definition':
+                node = node.named_children[0].next_named_sibling
             return node
     return None
 
@@ -278,4 +280,46 @@ def capture2identifier_boundaries(captures, lines: Sequence[str]) -> list[Identi
     unique_captures = {}
     for capture in captures:
         unique_captures[f'{capture.range[0]}:{capture.capture_type}'] = capture
-    return associate_identifier_parts(unique_captures.values(), lines)
+    # unique_captures={
+    # '157:function.decorator': CaptureInfo(capture_type='function.decorator', node=<Node type=decorator, start_point=(157, 4), end_point=(157, 17)>),
+    # '158:function.definition': CaptureInfo(capture_type='function.definition', node=<Node type=function_definition, start_point=(158, 4), end_point=(207, 19)>),
+    # '159:function.body': CaptureInfo(capture_type='function.body', node=<Node type=block, start_point=(159, 8), end_point=(207, 19)>)
+    # }
+    return associate_identifier_parts(sort_captures(unique_captures), lines)
+
+def parse_capture_key(key):
+    """
+    Parses the dictionary key into line number and capture type.
+    Args:
+        key (str): The key in the format 'line_number:capture_type'.
+    Returns:
+        tuple: (line_number as int, capture_type as str)
+    """
+    line_number, capture_type = key.split(':')
+    return int(line_number), capture_type.split('.')[-1]
+
+def get_sort_priority():
+    """
+    Returns a dictionary mapping capture types to their sort priority.
+    Returns:
+        dict: Capture type priorities.
+    """
+    return {'definition': 1, 'decorator': 2, 'body': 3, 'docstring': 4}
+
+def sort_captures(captures):
+    """
+    Sorts the values of the captures dictionary by capture type and line number.
+    Args:
+        captures (dict): The dictionary to sort.
+    Returns:
+        list: Sorted list of values.
+    """
+    priority = get_sort_priority()
+    sorted_items = sorted(
+        captures.items(),
+        key=lambda item: (
+            priority[parse_capture_key(item[0])[1]],  # Sort by capture type priority
+            parse_capture_key(item[0])[0]  # Then by line number
+        )
+    )
+    return [value for _, value in sorted_items]
diff --git a/src/cedarscript_editor/tree_sitter_identifier_queries.py b/src/cedarscript_editor/tree_sitter_identifier_queries.py
@@ -30,126 +30,110 @@
 #     except KeyError:
 #         return
 
+_common_template = """
+    ; Common pattern for body and docstring capture
+    body: (block
+        .
+        (expression_statement
+            (string) @{type}.docstring)?
+        .
+    ) @{type}.body
+"""
+
+_definition_base_template = """
+    name: (identifier) @_{type}_name
+    (#match? @_{type}_name "^{{name}}$")
+    (#set! role name)
+"""
 
 LANG_TO_TREE_SITTER_QUERY = {
     "python": {
         'function': """
-    ; Regular and async function definitions with optional docstring
-    (function_definition
-      name: (identifier) @_function_name
-      (#match? @_function_name "^{name}$")
-      body: (block) @function.body) @function.definition
-
+; Function Definitions
+(function_definition
+    {definition_base}
+    {common_body}
+) @function.definition
+
+; Decorated Function Definitions
+(decorated_definition
+    (decorator)+ @function.decorator
     (function_definition
-      name: (identifier) @_function_name
-      (#match? @_function_name "^{name}$")
-      body: (block
-        .
-        (expression_statement
-          (string) @function.docstring)?
-        .
-        (_)*))
-
-    ; Decorated function definitions (including async) with optional docstring
-    (decorated_definition
-      (decorator)+
-      (function_definition
-        name: (identifier) @_function_name
-        (#match? @_function_name "^{name}$")
-        body: (block) @function.body)) @function.definition
-
-    (decorated_definition
-      (decorator)+
-      (function_definition
-        name: (identifier) @_function_name
-        (#match? @_function_name "^{name}$")
-        body: (block
-          .
-          (expression_statement
-            (string) @function.docstring)?
-          .
-          (_)*)))
-
-    ; Method definitions in classes (including async and decorated) with optional docstring
-    (class_definition
-      body: (block
+        {definition_base}
+        {common_body}
+    ) @function.definition
+)
+
+; Methods in Classes
+(class_definition
+    body: (block
         (function_definition
-          name: (identifier) @_function_name
-          (#match? @_function_name "^{name}$")
-          body: (block) @function.body) @function.definition))
-
-    (class_definition
-      body: (block
-        (function_definition
-          name: (identifier) @_function_name
-          (#match? @_function_name "^{name}$")
-          body: (block
-            .
-            (expression_statement
-              (string) @function.docstring)?
-            .
-            (_)*))))
-""",
+            {definition_base}
+            {common_body}
+        ) @function.definition
+    )
+)
+
+; Decorated Methods in Classes
+(class_definition
+    body: (block
+        (decorated_definition
+            (decorator)+ @function.decorator
+            (function_definition
+                {definition_base}
+                {common_body}
+            ) @function.definition
+        )
+    )
+)
+""".format(
+            definition_base=_definition_base_template.format(type="function"),
+            common_body=_common_template.format(type="function")
+        ),
 
         'class': """
-    ; Regular and decorated class definitions (including nested) with optional docstring
-    (class_definition
-      name: (identifier) @_class_name
-      (#match? @_class_name "^{name}$")
-      body: (block) @class.body) @class.definition
-
+; Class Definitions
+(class_definition
+    {definition_base}
+    {common_body}
+) @class.definition
+
+; Decorated Class Definitions
+(decorated_definition
+    (decorator)+ @class.decorator
     (class_definition
-      name: (identifier) @_class_name
-      (#match? @_class_name "^{name}$")
-      body: (block
-        .
-        (expression_statement
-          (string) @class.docstring)?
-        .
-        (_)*))
-
-    ; Decorated class definitions
-    (decorated_definition
-      (decorator)+
-      (class_definition
-        name: (identifier) @_class_name
-        (#match? @_class_name "^{name}$")
-        body: (block) @class.body)) @class.definition
-
-    (decorated_definition
-      (decorator)+
-      (class_definition
-        name: (identifier) @_class_name
-        (#match? @_class_name "^{name}$")
-        body: (block
-          .
-          (expression_statement
-            (string) @class.docstring)?
-          .
-          (_)*)))
-
-    ; Nested class definitions within other classes
-    (class_definition
-      body: (block
-        (class_definition
-          name: (identifier) @_class_name
-          (#match? @_class_name "^{name}$")
-          body: (block) @class.body) @class.definition))
-
-    (class_definition
-      body: (block
+        {definition_base}
+        {common_body}
+    ) @class.definition
+)
+
+; Nested Classes
+(class_definition
+    body: (block
         (class_definition
-          name: (identifier) @_class_name
-          (#match? @_class_name "^{name}$")
-          body: (block
-            .
-            (expression_statement
-              (string) @class.docstring)?
-            .
-            (_)*))))
-"""
-    },
-    "kotlin": {
+            {definition_base}
+            {common_body}
+        ) @class.definition
+    )
+)
+
+; Decorated Nested Classes
+(class_definition
+    body: (block
+        (decorated_definition
+            (decorator)+ @class.decorator
+            (class_definition
+                {definition_base}
+                {common_body}
+            ) @class.definition
+        )
+    )
+)
+""".format(
+            definition_base=_definition_base_template.format(type="class"),
+            common_body=_common_template.format(type="class")
+        )
+    }, "kotlin": {
         'function': """
     ; Regular function definitions with optional annotations and KDoc
     (function_declaration
diff --git a/src/text_manipulation/range_spec.py b/src/text_manipulation/range_spec.py
@@ -13,6 +13,7 @@
 from collections.abc import Sequence
 from typing import NamedTuple, TypeAlias
 from functools import total_ordering
+from dataclasses import dataclass, field
 
 
 from cedarscript_ast_parser import Marker, RelativeMarker, RelativePositionType, MarkerType, BodyOrWhole
@@ -331,8 +332,8 @@ class ParentInfo(NamedTuple):
 
 ParentRestriction: TypeAlias = RangeSpec | str | None
 
-
-class IdentifierBoundaries(NamedTuple):
+@dataclass
+class IdentifierBoundaries:
     """
     Represents the boundaries of an identifier in code, including its whole range and body range.
 
@@ -347,8 +348,12 @@ class IdentifierBoundaries(NamedTuple):
     whole: RangeSpec
     body: RangeSpec | None = None
     docstring: RangeSpec | None = None
-    decorators: list[RangeSpec] = []
-    parents: list[ParentInfo] = []
+    decorators: list[RangeSpec] = field(default_factory=list)
+    parents: list[ParentInfo] = field(default_factory=list)
+
+    def append_decorator(self, decorator: RangeSpec):
+        self.decorators.append(decorator)
+        self.whole = self.whole._replace(start = min(self.whole.start, decorator.start))
 
     def __str__(self):
         return f'IdentifierBoundaries({self.whole} (BODY: {self.body}) )'
diff --git a/src/text_manipulation/text_editor_kit.py b/src/text_manipulation/text_editor_kit.py
@@ -8,7 +8,7 @@
 
 from collections.abc import Sequence
 from typing import Protocol, runtime_checkable
-from os import PathLike
+from os import PathLike, path
 
 from cedarscript_ast_parser import Marker, RelativeMarker, RelativePositionType, Segment, MarkerType, BodyOrWhole
 from .range_spec import IdentifierBoundaries, RangeSpec
@@ -24,7 +24,7 @@ def read_file(file_path: str | PathLike) -> str:
     Returns:
         str: The contents of the file as a string.
     """
-    with open(file_path, 'r') as file:
+    with open(path.normpath(file_path), 'r') as file:
         return file.read()
 
 
@@ -36,7 +36,7 @@ def write_file(file_path: str | PathLike, lines: Sequence[str]):
         file_path (str | PathLike): The path to the file to be written.
         lines (Sequence[str]): The lines to be written to the file.
     """
-    with open(file_path, 'w') as file:
+    with open(path.normpath(file_path), 'w') as file:
         file.writelines([line + '\n' for line in lines])
 
 
