From 421cfbe53bee1dcc1232dc3f71c756b13cb6d37c Mon Sep 17 00:00:00 2001
From: Daniel D'Avella <dan.davella@pixee.ai>
Date: Wed, 28 Feb 2024 21:09:28 -0500
Subject: [PATCH] Update result location matching; enable customizable logic

---
 src/codemodder/codemods/base_codemod.py |  4 ++-
 src/codemodder/result.py                | 45 +++++++++++++++++++------
 2 files changed, 37 insertions(+), 12 deletions(-)

diff --git a/src/codemodder/codemods/base_codemod.py b/src/codemodder/codemods/base_codemod.py
index a8e4d36f..ca3025ea 100644
--- a/src/codemodder/codemods/base_codemod.py
+++ b/src/codemodder/codemods/base_codemod.py
@@ -188,7 +188,7 @@ def _process_file(
             findings_for_rule = []
             for rule in rules:
                 findings_for_rule.extend(
-                    results.results_for_rule_and_file(rule, filename)
+                    results.results_for_rule_and_file(context, rule, filename)
                 )
 
         file_context = FileContext(
@@ -199,6 +199,8 @@ def _process_file(
             findings_for_rule,
         )
 
+        # TODO: for SAST tools we should preemtively filter out files that are not part of the result set
+
         if change_set := self.transformer.apply(
             context, file_context, findings_for_rule
         ):
diff --git a/src/codemodder/result.py b/src/codemodder/result.py
index 7e38f43d..8320b07d 100644
--- a/src/codemodder/result.py
+++ b/src/codemodder/result.py
@@ -1,9 +1,16 @@
+from __future__ import annotations
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Any
+from typing import Any, TYPE_CHECKING
+
+import libcst as cst
+from libcst._position import CodeRange
 
 from .utils.abc_dataclass import ABCDataclass
 
+if TYPE_CHECKING:
+    from codemodder.context import CodemodExecutionContext
+
 
 @dataclass
 class LineInfo:
@@ -24,16 +31,20 @@ class Result(ABCDataclass):
     rule_id: str
     locations: list[Location]
 
-    def match_location(self, pos, node):
-        for location in self.locations:
-            start_column = location.start.column
-            end_column = location.end.column
-            return (
-                pos.start.line == location.start.line
-                and (pos.start.column in (start_column - 1, start_column))
-                and pos.end.line == location.end.line
-                and (pos.end.column in (end_column - 1, end_column))
+    def match_location(self, pos: CodeRange, node: cst.CSTNode) -> bool:
+        del node
+        return any(
+            pos.start.line == location.start.line
+            and (
+                pos.start.column
+                in ((start_column := location.start.column) - 1, start_column)
+            )
+            and pos.end.line == location.end.line
+            and (
+                pos.end.column in ((end_column := location.end.column) - 1, end_column)
             )
+            for location in self.locations
+        )
 
 
 class ResultSet(dict[str, dict[Path, list[Result]]]):
@@ -41,7 +52,19 @@ def add_result(self, result: Result):
         for loc in result.locations:
             self.setdefault(result.rule_id, {}).setdefault(loc.file, []).append(result)
 
-    def results_for_rule_and_file(self, rule_id: str, file: Path) -> list[Result]:
+    def results_for_rule_and_file(
+        self, context: CodemodExecutionContext, rule_id: str, file: Path
+    ) -> list[Result]:
+        """
+        Return list of results for a given rule and file.
+
+        :param context: The codemod execution context
+        :param rule_id: The rule ID
+        :param file: The filename
+
+        Some implementers may need to use the context to compute paths that are relative to the target directory.
+        """
+        del context
         return self.get(rule_id, {}).get(file, [])
 
     def files_for_rule(self, rule_id: str) -> list[Path]: