Merge branch 'fix/java/line-profiler' of github.com:codeflash-ai/codeflash into fix/java/line-profiler

Author: HeshamHM28

codeflash/code_utils/tabulate.py

@@ -705,7 +705,7 @@ def tabulate(
     # format rows and columns, convert numeric values to strings
     cols = list(izip_longest(*list_of_lists))
     numparses = _expand_numparse(disable_numparse, len(cols))
-    coltypes = [_column_type(col, numparse=np) for col, np in zip(cols, numparses)]
+    coltypes = [_column_type(col, has_invisible, numparse=np) for col, np in zip(cols, numparses)]
     if isinstance(floatfmt, str):  # old version
         float_formats = len(cols) * [floatfmt]  # just duplicate the string to use in each column
     else:  # if floatfmt is list, tuple etc we have one per column

codeflash/languages/java/line_profiler.py

@@ -13,10 +13,9 @@
 
 import json
 import logging
-import os
 import re
 from pathlib import Path
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 if TYPE_CHECKING:
     from tree_sitter import Node
@@ -73,11 +72,7 @@ def __init__(self, output_file: Path, warmup_iterations: int = DEFAULT_WARMUP_IT
     # === Agent-based profiling (bytecode instrumentation) ===
 
     def generate_agent_config(
-        self,
-        source: str,
-        file_path: Path,
-        functions: list[FunctionInfo],
-        config_output_path: Path,
+        self, source: str, file_path: Path, functions: list[FunctionInfo], config_output_path: Path
     ) -> Path:
         """Generate config JSON for the profiler agent.
 
@@ -104,7 +99,12 @@ class name, then writes a config JSON that the agent uses to know which
             for line_num in range(func.starting_line, func.ending_line + 1):
                 if 1 <= line_num <= len(lines):
                     content = lines[line_num - 1].strip()
-                    if content and not content.startswith("//") and not content.startswith("/*") and not content.startswith("*"):
+                    if (
+                        content
+                        and not content.startswith("//")
+                        and not content.startswith("/*")
+                        and not content.startswith("*")
+                    ):
                         key = f"{file_path.as_posix()}:{line_num}"
                         line_contents[key] = content
 
@@ -120,12 +120,7 @@ class name, then writes a config JSON that the agent uses to know which
         config = {
             "outputFile": str(self.output_file),
             "warmupIterations": self.warmup_iterations,
-            "targets": [
-                {
-                    "className": class_name,
-                    "methods": method_targets,
-                }
-            ],
+            "targets": [{"className": class_name, "methods": method_targets}],
             "lineContents": line_contents,
         }
 
@@ -143,7 +138,9 @@ def build_javaagent_arg(self, config_path: Path) -> str:
 
     # === Source-level instrumentation ===
 
-    def instrument_source(self, source: str, file_path: Path, functions: list[FunctionInfo], analyzer=None) -> str:
+    def instrument_source(
+        self, source: str, file_path: Path, functions: list[FunctionInfo], analyzer: Any = None
+    ) -> str:
         """Instrument Java source code with line profiling.
 
         Injects a profiler class and per-line hit() calls directly into the source.
@@ -159,7 +156,7 @@ def instrument_source(self, source: str, file_path: Path, functions: list[Functi
 
         """
         # Initialize line contents map
-        self.line_contents = {}
+        self.line_contents: dict[str, str] = {}
 
         lines = source.splitlines(keepends=True)
 
@@ -327,7 +324,7 @@ class {self.profiler_class} {{
 }}
 """
 
-    def instrument_function(self, func: FunctionInfo, lines: list[str], file_path: Path, analyzer) -> list[str]:
+    def instrument_function(self, func: FunctionInfo, lines: list[str], file_path: Path, analyzer: Any) -> list[str]:
         """Instrument a single function with line profiling.
 
         Args:
@@ -446,7 +443,7 @@ def find_executable_lines(self, node: Node) -> set[int]:
     # === Result parsing (shared by both approaches) ===
 
     @staticmethod
-    def parse_results(profile_file: Path) -> dict:
+    def parse_results(profile_file: Path) -> dict[str, Any]:
         """Parse line profiling results from the agent's JSON output.
 
         Returns the same format as parse_line_profile_test_output.parse_line_profile_results()
@@ -520,9 +517,9 @@ def parse_results(profile_file: Path) -> dict:
                 for fp, line_stats in lines_by_file.items():
                     sorted_stats = sorted(line_stats, key=lambda t: t[0])
                     if sorted_stats:
-                        grouped_timings[(fp, sorted_stats[0][0], os.path.basename(fp))] = sorted_stats
+                        grouped_timings[(fp, sorted_stats[0][0], Path(fp).name)] = sorted_stats
 
-            result: dict = {"timings": grouped_timings, "unit": 1e-9, "line_contents": line_contents}
+            result: dict[str, Any] = {"timings": grouped_timings, "unit": 1e-9, "line_contents": line_contents}
             result["str_out"] = format_line_profile_results(result, line_contents)
             return result
 
@@ -531,15 +528,14 @@ def parse_results(profile_file: Path) -> dict:
             return {"timings": {}, "unit": 1e-9, "str_out": ""}
 
 
-def load_method_ranges(
-    profile_file: Path,
-) -> tuple[list[tuple[str, str, int, int]], dict[str, str]]:
+def load_method_ranges(profile_file: Path) -> tuple[list[tuple[str, str, int, int]], dict[str, str]]:
     """Load method ranges and line contents from the agent config file.
 
     Returns:
         (method_ranges, config_line_contents) where method_ranges is a list of
         (source_file, method_name, start_line, end_line) and config_line_contents
         is the lineContents dict from the config (key: "file:line", value: source text).
+
     """
     config_path = profile_file.with_suffix(".config.json")
     if not config_path.exists():
@@ -549,12 +545,7 @@ def load_method_ranges(
         ranges = []
         for target in config.get("targets", []):
             for method in target.get("methods", []):
-                ranges.append((
-                    method.get("sourceFile", ""),
-                    method["name"],
-                    method["startLine"],
-                    method["endLine"],
-                ))
+                ranges.append((method.get("sourceFile", ""), method["name"], method["startLine"], method["endLine"]))
         return ranges, config.get("lineContents", {})
     except Exception:
         return [], {}
@@ -571,7 +562,7 @@ def find_method_for_line(
     for source_file, method_name, start_line, end_line in method_ranges:
         if file_path == source_file and start_line <= line_num <= end_line:
             return method_name, start_line
-    return os.path.basename(file_path), line_num
+    return Path(file_path).name, line_num
 
 
 def find_agent_jar() -> Path | None:
@@ -580,16 +571,7 @@ def find_agent_jar() -> Path | None:
     Checks local Maven repo, package resources, and development build directory.
     """
     # Check local Maven repository first (fastest)
-    m2_jar = (
-        Path.home()
-        / ".m2"
-        / "repository"
-        / "com"
-        / "codeflash"
-        / "codeflash-runtime"
-        / "1.0.0"
-        / AGENT_JAR_NAME
-    )
+    m2_jar = Path.home() / ".m2" / "repository" / "com" / "codeflash" / "codeflash-runtime" / "1.0.0" / AGENT_JAR_NAME
     if m2_jar.exists():
         return m2_jar
 
@@ -599,9 +581,7 @@ def find_agent_jar() -> Path | None:
         return resources_jar
 
     # Check development build directory
-    dev_jar = (
-        Path(__file__).parent.parent.parent.parent / "codeflash-java-runtime" / "target" / AGENT_JAR_NAME
-    )
+    dev_jar = Path(__file__).parent.parent.parent.parent / "codeflash-java-runtime" / "target" / AGENT_JAR_NAME
     if dev_jar.exists():
         return dev_jar
 
@@ -624,7 +604,7 @@ def resolve_internal_class_name(file_path: Path, source: str) -> str:
 
 
 def format_line_profile_results(
-    results: dict, line_contents: dict[tuple[str, int], str] | None = None
+    results: dict[str, Any], line_contents: dict[tuple[str, int], str] | None = None
 ) -> str:
     """Format line profiling results using the same tabulate pipe format as Python.
 

codeflash/languages/java/support.py

@@ -273,7 +273,7 @@ def _build_runtime_map(self, inv_id_runtimes: dict[InvocationId, list[int]]) ->
 
     def compare_test_results(
         self, original_results_path: Path, candidate_results_path: Path, project_root: Path | None = None
-    ) -> tuple[bool, list]:
+    ) -> tuple[bool, list[Any]]:
         """Compare test results between original and candidate code."""
         return _compare_test_results(original_results_path, candidate_results_path, project_root=project_root)
 
@@ -409,10 +409,7 @@ def instrument_source_for_line_profiler(
 
             config_path = line_profiler_output_file.with_suffix(".config.json")
             profiler.generate_agent_config(
-                source=source,
-                file_path=func_info.file_path,
-                functions=[func_info],
-                config_output_path=config_path,
+                source=source, file_path=func_info.file_path, functions=[func_info], config_output_path=config_path
             )
 
             self.line_profiler_agent_arg = profiler.build_javaagent_arg(config_path)
@@ -422,7 +419,7 @@ def instrument_source_for_line_profiler(
             logger.exception("Failed to prepare line profiling for %s", func_info.function_name)
             return False
 
-    def parse_line_profile_results(self, line_profiler_output_file: Path) -> dict:
+    def parse_line_profile_results(self, line_profiler_output_file: Path) -> dict[str, Any]:
         """Parse line profiler output for Java.
 
         Args:

codeflash/languages/registry.py

@@ -58,9 +58,6 @@ def _ensure_languages_registered() -> None:
     with contextlib.suppress(ImportError):
         from codeflash.languages.java import support as _  # noqa: F401
 
-    with contextlib.suppress(ImportError):
-        from codeflash.languages.java import support as _  # noqa: F401
-
     _languages_registered = True