From 2f6a91f95d84baab90ddfae09b4b324a24a8ee55 Mon Sep 17 00:00:00 2001
From: Sarthak Agarwal <sarthak.saga@gmail.com>
Date: Thu, 29 Jan 2026 19:26:50 +0530
Subject: [PATCH] fix: lint issues in experiments folder and format fixes

---
 codeflash/api/aiservice.py                        | 4 +---
 codeflash/languages/base.py                       | 2 --
 experiments/__init__.py                           | 0
 experiments/code_replacement/__init__.py          | 0
 experiments/code_replacement/approach_c_hybrid.py | 8 +++++---
 experiments/code_replacement/run_experiments.py   | 4 ++--
 experiments/code_replacement/test_cases.py        | 2 +-
 7 files changed, 9 insertions(+), 11 deletions(-)
 create mode 100644 experiments/__init__.py
 create mode 100644 experiments/code_replacement/__init__.py

diff --git a/codeflash/api/aiservice.py b/codeflash/api/aiservice.py
index a60d34f21..364dccf8c 100644
--- a/codeflash/api/aiservice.py
+++ b/codeflash/api/aiservice.py
@@ -240,9 +240,7 @@ class AiServiceClient:
             n_candidates=n_candidates,
         )
 
-    def get_jit_rewritten_code(
-        self, source_code: str, trace_id: str
-    ) -> list[OptimizedCandidate]:
+    def get_jit_rewritten_code(self, source_code: str, trace_id: str) -> list[OptimizedCandidate]:
         """Rewrite the given python code for performance via jit compilation by making a request to the Django endpoint.
 
         Parameters
diff --git a/codeflash/languages/base.py b/codeflash/languages/base.py
index 965258b5f..11b5afd4f 100644
--- a/codeflash/languages/base.py
+++ b/codeflash/languages/base.py
@@ -613,8 +613,6 @@ class LanguageSupport(Protocol):
         """Parse line profiler output."""
         ...
 
-
-
     # === Test Execution ===
 
     def run_behavioral_tests(
diff --git a/experiments/__init__.py b/experiments/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/experiments/code_replacement/__init__.py b/experiments/code_replacement/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/experiments/code_replacement/approach_c_hybrid.py b/experiments/code_replacement/approach_c_hybrid.py
index dcf41353d..186e6af97 100644
--- a/experiments/code_replacement/approach_c_hybrid.py
+++ b/experiments/code_replacement/approach_c_hybrid.py
@@ -18,6 +18,7 @@ Cons:
 - Still text-based replacement (not AST rewriting)
 """
 
+import sys
 from dataclasses import dataclass
 from typing import Optional
 
@@ -117,7 +118,7 @@ class HybridReplacer:
 
             return None
 
-        def traverse(node):
+        def traverse(node) -> None:  # noqa: ANN001
             """Recursively traverse tree to find functions."""
             node_type = node.type
 
@@ -246,7 +247,8 @@ class HybridReplacer:
         boundaries = self.find_function_boundaries(source, function_name)
 
         if not boundaries:
-            raise ValueError(f"Function '{function_name}' not found in source")
+            msg = f"Function '{function_name}' not found in source"
+            raise ValueError(msg)
 
         if len(boundaries) > 1:
             # Multiple functions with same name - use the first one
@@ -368,7 +370,7 @@ if __name__ == "__main__":
 
     if not TREE_SITTER_AVAILABLE:
         print("Cannot run tests: tree-sitter not installed")
-        exit(1)
+        sys.exit(1)
 
     replacer = HybridReplacer("javascript")
     ts_replacer = HybridReplacer("typescript")
diff --git a/experiments/code_replacement/run_experiments.py b/experiments/code_replacement/run_experiments.py
index 78ca792cb..a312fae0f 100644
--- a/experiments/code_replacement/run_experiments.py
+++ b/experiments/code_replacement/run_experiments.py
@@ -159,7 +159,7 @@ def test_approach_a() -> ApproachSummary:
 
         replacer = JsCodeshiftReplacer()
 
-        if not replacer._check_node_available():
+        if not replacer._check_node_available():  # noqa: SLF001
             summary.available = False
             return summary
 
@@ -235,7 +235,7 @@ def generate_report(summaries: list[ApproachSummary]) -> str:
     return "\n".join(report)
 
 
-def main():
+def main() -> None:
     """Run all experiments and generate report."""
     print("=" * 70)
     print("Code Replacement Strategy Experiments")
diff --git a/experiments/code_replacement/test_cases.py b/experiments/code_replacement/test_cases.py
index 80d84718e..bfbd0992a 100644
--- a/experiments/code_replacement/test_cases.py
+++ b/experiments/code_replacement/test_cases.py
@@ -634,7 +634,7 @@ function last() {
 ]
 
 
-def get_test_cases():
+def get_test_cases() -> list[ReplacementTestCase]:
     """Return all test cases."""
     return TEST_CASES