refactor: delegate PythonSupport context methods to canonical pipeline

Replace duplicate implementations in extract_code_context() and
find_helper_functions() with calls to get_code_optimization_context()
and get_function_sources_from_jedi() from the canonical context module.

Author: KRRT7

codeflash/languages/python/support.py

@@ -171,127 +171,59 @@ def discover_tests(
     # === Code Analysis ===
 
     def extract_code_context(self, function: FunctionToOptimize, project_root: Path, module_root: Path) -> CodeContext:
-        """Extract function code and its dependencies.
+        """Extract function code and its dependencies via the canonical context pipeline."""
+        from codeflash.languages.python.context.code_context_extractor import get_code_optimization_context
 
-        Uses jedi and libcst for Python code analysis.
-
-        Args:
-            function: The function to extract context for.
-            project_root: Root of the project.
-            module_root: Root of the module containing the function.
-
-        Returns:
-            CodeContext with target code and dependencies.
-
-        """
         try:
-            source = function.file_path.read_text()
+            result = get_code_optimization_context(function, project_root)
         except Exception as e:
-            logger.exception("Failed to read %s: %s", function.file_path, e)
+            logger.warning("Failed to extract code context for %s: %s", function.function_name, e)
             return CodeContext(target_code="", target_file=function.file_path, language=Language.PYTHON)
 
-        # Extract the function source
-        lines = source.splitlines(keepends=True)
-        if function.starting_line and function.ending_line:
-            target_lines = lines[function.starting_line - 1 : function.ending_line]
-            target_code = "".join(target_lines)
-        else:
-            target_code = ""
-
-        # Find helper functions
-        helpers = self.find_helper_functions(function, project_root)
-
-        # Extract imports
-        import_lines = []
-        for line in lines:
-            stripped = line.strip()
-            if stripped.startswith(("import ", "from ")):
-                import_lines.append(stripped)
-            elif stripped and not stripped.startswith("#"):
-                # Stop at first non-import, non-comment line
-                break
+        helpers = [
+            HelperFunction(
+                name=fs.only_function_name,
+                qualified_name=fs.qualified_name,
+                file_path=fs.file_path,
+                source_code=fs.source_code,
+                start_line=fs.jedi_definition.line if fs.jedi_definition else 1,
+                end_line=fs.jedi_definition.line if fs.jedi_definition else 1,
+            )
+            for fs in result.helper_functions
+        ]
 
         return CodeContext(
-            target_code=target_code,
+            target_code=result.read_writable_code.markdown,
             target_file=function.file_path,
             helper_functions=helpers,
-            read_only_context="",
-            imports=import_lines,
+            read_only_context=result.read_only_context_code,
+            imports=[],
             language=Language.PYTHON,
         )
 
     def find_helper_functions(self, function: FunctionToOptimize, project_root: Path) -> list[HelperFunction]:
-        """Find helper functions called by the target function.
-
-        Uses jedi for Python code analysis.
-
-        Args:
-            function: The target function to analyze.
-            project_root: Root of the project.
-
-        Returns:
-            List of HelperFunction objects.
-
-        """
-        helpers: list[HelperFunction] = []
+        """Find helper functions called by the target function via the canonical jedi pipeline."""
+        from codeflash.languages.python.context.code_context_extractor import get_function_sources_from_jedi
 
         try:
-            import jedi
-
-            from codeflash.code_utils.code_utils import get_qualified_name, path_belongs_to_site_packages
-            from codeflash.optimization.function_context import belongs_to_function_qualified
-
-            script = jedi.Script(path=function.file_path, project=jedi.Project(path=project_root))
-            file_refs = script.get_names(all_scopes=True, definitions=False, references=True)
-
-            qualified_name = function.qualified_name
-
-            for ref in file_refs:
-                if not ref.full_name or not belongs_to_function_qualified(ref, qualified_name):
-                    continue
-
-                try:
-                    definitions = ref.goto(follow_imports=True, follow_builtin_imports=False)
-                except Exception:
-                    continue
-
-                for definition in definitions:
-                    definition_path = definition.module_path
-                    if definition_path is None:
-                        continue
-
-                    # Check if it's a valid helper (in project, not in target function)
-                    is_valid = (
-                        str(definition_path).startswith(str(project_root))
-                        and not path_belongs_to_site_packages(definition_path)
-                        and definition.full_name
-                        and not belongs_to_function_qualified(definition, qualified_name)
-                        and definition.type == "function"
-                    )
-
-                    if is_valid:
-                        helper_qualified_name = get_qualified_name(definition.module_name, definition.full_name)
-                        # Get source code
-                        try:
-                            helper_source = definition.get_line_code()
-                        except Exception:
-                            helper_source = ""
-
-                        helpers.append(
-                            HelperFunction(
-                                name=definition.name,
-                                qualified_name=helper_qualified_name,
-                                file_path=definition_path,
-                                source_code=helper_source,
-                                start_line=definition.line or 1,
-                                end_line=definition.line or 1,
-                            )
-                        )
-
+            _dict, sources = get_function_sources_from_jedi(
+                {function.file_path: {function.qualified_name}}, project_root
+            )
         except Exception as e:
             logger.warning("Failed to find helpers for %s: %s", function.function_name, e)
+            return []
 
-        return helpers
+        return [
+            HelperFunction(
+                name=fs.only_function_name,
+                qualified_name=fs.qualified_name,
+                file_path=fs.file_path,
+                source_code=fs.source_code,
+                start_line=fs.jedi_definition.line if fs.jedi_definition else 1,
+                end_line=fs.jedi_definition.line if fs.jedi_definition else 1,
+            )
+            for fs in sources
+        ]
 
     def find_references(
         self, function: FunctionToOptimize, project_root: Path, tests_root: Path | None = None, max_files: int = 500