refactors; adapt with filtering logic

Author: themisvaltinos

sqlmesh/lsp/reference.py

@@ -450,75 +450,76 @@ def get_model_find_all_references(
     Returns:
         A list of references to the model across all files
     """
-    # First, get the references in the current file to determine what model we're looking for
-    current_file_references = [
-        ref
-        for ref in get_model_definitions_for_a_path(lint_context, document_uri)
-        if isinstance(ref, LSPModelReference)
-    ]
-
     # Find the model reference at the cursor position
-    target_model_uri: t.Optional[str] = None
-    for ref in current_file_references:
-        if _position_within_range(position, ref.range):
-            # This is a model reference, get the target model URI
-            target_model_uri = ref.uri
-            break
+    model_at_position = next(
+        filter(
+            lambda ref: isinstance(ref, LSPModelReference)
+            and _position_within_range(position, ref.range),
+            get_model_definitions_for_a_path(lint_context, document_uri),
+        ),
+        None,
+    )
 
-    if target_model_uri is None:
+    if not model_at_position:
         return []
 
+    assert isinstance(model_at_position, LSPModelReference)  # for mypy
+
+    target_model_uri = model_at_position.uri
+
     # Start with the model definition
     all_references: t.List[LSPModelReference] = [
         LSPModelReference(
-            uri=ref.uri,
+            uri=model_at_position.uri,
             range=Range(
                 start=Position(line=0, character=0),
                 end=Position(line=0, character=0),
             ),
-            markdown_description=ref.markdown_description,
+            markdown_description=model_at_position.markdown_description,
         )
     ]
 
-    # Then add the original reference
-    for ref in current_file_references:
-        if ref.uri == target_model_uri and isinstance(ref, LSPModelReference):
-            all_references.append(
-                LSPModelReference(
-                    uri=document_uri.value,
-                    range=ref.range,
-                    markdown_description=ref.markdown_description,
-                )
+    # Then add references from the current file
+    current_file_refs = filter(
+        lambda ref: isinstance(ref, LSPModelReference) and ref.uri == target_model_uri,
+        get_model_definitions_for_a_path(lint_context, document_uri),
+    )
+
+    for ref in current_file_refs:
+        assert isinstance(ref, LSPModelReference)  # for mypy
+
+        all_references.append(
+            LSPModelReference(
+                uri=document_uri.value,
+                range=ref.range,
+                markdown_description=ref.markdown_description,
             )
+        )
 
     # Search through the models in the project
-    for path, target in lint_context.map.items():
-        if not isinstance(target, (ModelTarget, AuditTarget)):
-            continue
-
+    for path, _ in lint_context.map.items():
         file_uri = URI.from_path(path)
 
         # Skip current file, already processed
         if file_uri.value == document_uri.value:
             continue
 
-        # Get model references for this file
-        file_references = [
-            ref
-            for ref in get_model_definitions_for_a_path(lint_context, file_uri)
-            if isinstance(ref, LSPModelReference)
-        ]
-
-        # Add references that point to the target model file
-        for ref in file_references:
-            if ref.uri == target_model_uri and isinstance(ref, LSPModelReference):
-                all_references.append(
-                    LSPModelReference(
-                        uri=file_uri.value,
-                        range=ref.range,
-                        markdown_description=ref.markdown_description,
-                    )
+        # Get model references that point to the target model
+        matching_refs = filter(
+            lambda ref: isinstance(ref, LSPModelReference) and ref.uri == target_model_uri,
+            get_model_definitions_for_a_path(lint_context, file_uri),
+        )
+
+        for ref in matching_refs:
+            assert isinstance(ref, LSPModelReference)  # for mypy
+
+            all_references.append(
+                LSPModelReference(
+                    uri=file_uri.value,
+                    range=ref.range,
+                    markdown_description=ref.markdown_description,
                 )
+            )
 
     return all_references
 
@@ -588,41 +589,39 @@ def get_cte_references(
 
 
 def get_macro_find_all_references(
-    lint_context: LSPContext, document_uri: URI, position: Position
+    lsp_context: LSPContext, document_uri: URI, position: Position
 ) -> t.List[LSPMacroReference]:
     """
     Get all references to a macro at a specific position in a document.
 
     This function finds all usages of a macro across the entire project.
 
     Args:
-        lint_context: The LSP context
+        lsp_context: The LSP context
         document_uri: The URI of the document
         position: The position to check for macro references
 
     Returns:
         A list of references to the macro across all files
     """
-    # First, get the macro references in the current file to identify the target macro
-    current_file_references = [
-        ref
-        for ref in get_macro_definitions_for_a_path(lint_context, document_uri)
-        if isinstance(ref, LSPMacroReference)
-    ]
-
     # Find the macro reference at the cursor position
-    target_macro_uri: t.Optional[str] = None
-    target_macro_target_range: t.Optional[Range] = None
-
-    for ref in current_file_references:
-        if _position_within_range(position, ref.range):
-            target_macro_uri = ref.uri
-            target_macro_target_range = ref.target_range
-            break
+    macro_at_position = next(
+        filter(
+            lambda ref: isinstance(ref, LSPMacroReference)
+            and _position_within_range(position, ref.range),
+            get_macro_definitions_for_a_path(lsp_context, document_uri),
+        ),
+        None,
+    )
 
-    if target_macro_uri is None:
+    if not macro_at_position:
         return []
 
+    assert isinstance(macro_at_position, LSPMacroReference)  # for mypy
+
+    target_macro_uri = macro_at_position.uri
+    target_macro_target_range = macro_at_position.target_range
+
     # Start with the macro definition
     all_references: t.List[LSPMacroReference] = [
         LSPMacroReference(
@@ -634,30 +633,27 @@ def get_macro_find_all_references(
     ]
 
     # Search through all SQL and audit files in the project
-    for path, target in lint_context.map.items():
-        if not isinstance(target, (ModelTarget, AuditTarget)):
-            continue
-
+    for path, _ in lsp_context.map.items():
         file_uri = URI.from_path(path)
 
-        # Get macro references for this file
-        file_macro_references = [
-            ref
-            for ref in get_macro_definitions_for_a_path(lint_context, file_uri)
-            if isinstance(ref, LSPMacroReference)
-        ]
-
-        # Add references that point to the same macro definition
-        for ref in file_macro_references:
-            if ref.uri == target_macro_uri and ref.target_range == target_macro_target_range:
-                all_references.append(
-                    LSPMacroReference(
-                        uri=file_uri.value,
-                        range=ref.range,
-                        target_range=ref.target_range,
-                        markdown_description=ref.markdown_description,
-                    )
+        # Get macro references that point to the same macro definition
+        matching_refs = filter(
+            lambda ref: isinstance(ref, LSPMacroReference)
+            and ref.uri == target_macro_uri
+            and ref.target_range == target_macro_target_range,
+            get_macro_definitions_for_a_path(lsp_context, file_uri),
+        )
+
+        for ref in matching_refs:
+            assert isinstance(ref, LSPMacroReference)  # for mypy
+            all_references.append(
+                LSPMacroReference(
+                    uri=file_uri.value,
+                    range=ref.range,
+                    target_range=ref.target_range,
+                    markdown_description=ref.markdown_description,
                 )
+            )
 
     return all_references
 

tests/lsp/test_reference_macro_find_all.py

@@ -6,6 +6,11 @@
     get_macro_definitions_for_a_path,
 )
 from sqlmesh.lsp.uri import URI
+from sqlmesh.core.linter.helpers import (
+    read_range_from_file,
+    Range as SQLMeshRange,
+    Position as SQLMeshPosition,
+)
 
 
 def test_find_all_references_for_macro_add_one():
@@ -43,35 +48,32 @@ def test_find_all_references_for_macro_add_one():
     usage_refs = [ref for ref in all_references if "top_waiters" in ref.uri]
     assert len(usage_refs) >= 1, "Should include the usage in top_waiters.sql"
 
-    # breakpoint()
-    expected_ranges = [
-        # Macro definition in utils.py
-        {
-            "uri_file": "utils.py",
-            "range": ((6, 0), (9, 22)),
-        },
-        # Usage in customers.sql
-        {
-            "uri_file": "customers.sql",
-            "range": ((36, 7), (36, 14)),
-        },
-        # Usage in top_waiters.sql
-        {
-            "uri_file": "top_waiters.sql",
-            "range": ((12, 5), (12, 12)),
-        },
-    ]
-
-    for expected in expected_ranges:
-        assert any(
-            expected["uri_file"] in ref.uri
-            and ref.range.start.line == expected["range"][0][0]
-            and ref.range.start.character == expected["range"][0][1]
-            and ref.range.end.line == expected["range"][1][0]
-            and ref.range.end.character == expected["range"][1][1]
-            for ref in all_references
-        ), (
-            f"Expected reference with uri {expected['uri_file']} and range {expected['range']} not found"
+    expected_files = {
+        "utils.py": {"pattern": r"def add_one", "expected_content": "def add_one"},
+        "customers.sql": {"pattern": r"@ADD_ONE\s*\(", "expected_content": "ADD_ONE"},
+        "top_waiters.sql": {"pattern": r"@ADD_ONE\s*\(", "expected_content": "ADD_ONE"},
+    }
+
+    for expected_file, expectations in expected_files.items():
+        file_refs = [ref for ref in all_references if expected_file in ref.uri]
+        assert len(file_refs) >= 1, f"Should find at least one reference in {expected_file}"
+
+        file_ref = file_refs[0]
+        file_path = URI(file_ref.uri).to_path()
+
+        sqlmesh_range = SQLMeshRange(
+            start=SQLMeshPosition(
+                line=file_ref.range.start.line, character=file_ref.range.start.character
+            ),
+            end=SQLMeshPosition(
+                line=file_ref.range.end.line, character=file_ref.range.end.character
+            ),
+        )
+
+        # Read the content at the reference location
+        content = read_range_from_file(file_path, sqlmesh_range)
+        assert content.startswith(expectations["expected_content"]), (
+            f"Expected content to start with '{expectations['expected_content']}', got: {content}"
         )