global language config

Author: misrasaurabh1

codeflash/api/aiservice.py

@@ -14,6 +14,7 @@
 from codeflash.code_utils.env_utils import get_codeflash_api_key
 from codeflash.code_utils.git_utils import get_last_commit_author_if_pr_exists, get_repo_owner_and_name
 from codeflash.code_utils.time_utils import humanize_runtime
+from codeflash.languages import is_javascript, is_python
 from codeflash.models.ExperimentMetadata import ExperimentMetadata
 from codeflash.models.models import (
     AIServiceRefinerRequest,
@@ -178,7 +179,7 @@ def optimize_code(  # noqa: D417
         # Add language-specific version fields
         # Always include python_version for backward compatibility with older backend
         payload["python_version"] = platform.python_version()
-        if language == "python":
+        if is_python():
             pass  # python_version already set
         else:
             payload["language_version"] = language_version or "ES2022"
@@ -432,7 +433,7 @@ def optimize_code_refinement(self, request: list[AIServiceRefinerRequest]) -> li
 
             # Add language version - always include python_version for backward compatibility
             item["python_version"] = platform.python_version()
-            if opt.language == "python":
+            if is_python():
                 pass  # python_version already set
             elif opt.language_version:
                 item["language_version"] = opt.language_version
@@ -752,11 +753,11 @@ def generate_regression_tests(  # noqa: D417
         # Validate test framework based on language
         python_frameworks = ["pytest", "unittest"]
         javascript_frameworks = ["jest", "mocha", "vitest"]
-        if language == "python":
+        if is_python():
             assert test_framework in python_frameworks, (
                 f"Invalid test framework for Python, got {test_framework} but expected one of {python_frameworks}"
             )
-        elif language in ("javascript", "typescript"):
+        elif is_javascript():
             assert test_framework in javascript_frameworks, (
                 f"Invalid test framework for JavaScript, got {test_framework} but expected one of {javascript_frameworks}"
             )
@@ -781,7 +782,7 @@ def generate_regression_tests(  # noqa: D417
         # Add language-specific version fields
         # Always include python_version for backward compatibility with older backend
         payload["python_version"] = platform.python_version()
-        if language == "python":
+        if is_python():
             pass  # python_version already set
         else:
             payload["language_version"] = language_version or "ES2022"
@@ -873,7 +874,7 @@ def get_optimization_review(
             "codeflash_version": codeflash_version,
             "calling_fn_details": calling_fn_details,
             "language": language,
-            "python_version": platform.python_version() if language == "python" else None,
+            "python_version": platform.python_version() if is_python() else None,
             "call_sequence": self.get_next_sequence(),
         }
         console.rule()

codeflash/code_utils/code_replacer.py

@@ -18,6 +18,7 @@
 from codeflash.code_utils.config_parser import find_conftest_files
 from codeflash.code_utils.formatter import sort_imports
 from codeflash.code_utils.line_profile_utils import ImportAdder
+from codeflash.languages import is_python
 from codeflash.models.models import FunctionParent
 
 if TYPE_CHECKING:
@@ -445,7 +446,7 @@ def replace_function_definitions_in_module(
     function_to_optimize: Optional[FunctionToOptimize] = None,
 ) -> bool:
     # Route to language-specific implementation for non-Python languages
-    if optimized_code.language and optimized_code.language != "python":
+    if not is_python():
         return replace_function_definitions_for_language(
             function_names, optimized_code, module_abspath, project_root_path, function_to_optimize
         )

codeflash/code_utils/deduplicate_code.py

@@ -11,6 +11,7 @@
 from typing import TYPE_CHECKING
 
 from codeflash.code_utils.normalizers import get_normalizer
+from codeflash.languages import current_language, is_python
 
 if TYPE_CHECKING:
     pass
@@ -20,7 +21,7 @@ def normalize_code(
     code: str,
     remove_docstrings: bool = True,  # noqa: FBT001, FBT002
     return_ast_dump: bool = False,  # noqa: FBT001, FBT002
-    language: str = "python",
+    language: str | None = None,
 ) -> str:
     """Normalize code by parsing, cleaning, and normalizing variable names.
 
@@ -30,16 +31,19 @@ def normalize_code(
         code: Source code as string
         remove_docstrings: Whether to remove docstrings (Python only)
         return_ast_dump: Return AST dump instead of unparsed code (Python only)
-        language: Language of the code ('python', 'javascript', 'typescript')
+        language: Language of the code. If None, uses the current session language.
 
     Returns:
         Normalized code as string
     """
+    if language is None:
+        language = current_language().value
+
     try:
         normalizer = get_normalizer(language)
 
         # Python has additional options
-        if language == "python":
+        if is_python():
             if return_ast_dump:
                 return normalizer.normalize_for_hash(code)
             return normalizer.normalize(code, remove_docstrings=remove_docstrings)
@@ -51,7 +55,7 @@ def normalize_code(
         return _basic_normalize(code)
     except Exception:
         # Parsing error - try other languages or fall back
-        if language == "python":
+        if is_python():
             # Try JavaScript as fallback
             try:
                 js_normalizer = get_normalizer("javascript")
@@ -76,16 +80,19 @@ def _basic_normalize(code: str) -> str:
     return " ".join(code.split())
 
 
-def get_code_fingerprint(code: str, language: str = "python") -> str:
+def get_code_fingerprint(code: str, language: str | None = None) -> str:
     """Generate a fingerprint for normalized code.
 
     Args:
         code: Source code
-        language: Language of the code ('python', 'javascript', 'typescript')
+        language: Language of the code. If None, uses the current session language.
 
     Returns:
         SHA-256 hash of normalized code
     """
+    if language is None:
+        language = current_language().value
+
     try:
         normalizer = get_normalizer(language)
         return normalizer.get_fingerprint(code)
@@ -95,17 +102,20 @@ def get_code_fingerprint(code: str, language: str = "python") -> str:
         return hashlib.sha256(normalized.encode()).hexdigest()
 
 
-def are_codes_duplicate(code1: str, code2: str, language: str = "python") -> bool:
+def are_codes_duplicate(code1: str, code2: str, language: str | None = None) -> bool:
     """Check if two code segments are duplicates after normalization.
 
     Args:
         code1: First code segment
         code2: Second code segment
-        language: Language of the code ('python', 'javascript', 'typescript')
+        language: Language of the code. If None, uses the current session language.
 
     Returns:
         True if codes are structurally identical (ignoring local variable names)
     """
+    if language is None:
+        language = current_language().value
+
     try:
         normalizer = get_normalizer(language)
         return normalizer.are_duplicates(code1, code2)

codeflash/context/code_context_extractor.py

@@ -23,6 +23,7 @@
 from codeflash.discovery.functions_to_optimize import FunctionToOptimize  # noqa: TC001
 
 # Language support imports for multi-language code context extraction
+from codeflash.languages import is_python
 from codeflash.languages.base import Language
 from codeflash.models.models import (
     CodeContextType,
@@ -79,7 +80,7 @@ def get_code_optimization_context(
     testgen_token_limit: int = TESTGEN_CONTEXT_TOKEN_LIMIT,
 ) -> CodeOptimizationContext:
     # Route to language-specific implementation for non-Python languages
-    if function_to_optimize.language and function_to_optimize.language != "python":
+    if not is_python():
         return get_code_optimization_context_for_language(
             function_to_optimize, project_root_path, optim_token_limit, testgen_token_limit
         )

codeflash/context/unused_definition_remover.py

@@ -11,6 +11,7 @@
 
 from codeflash.cli_cmds.console import logger
 from codeflash.code_utils.code_replacer import replace_function_definitions_in_module
+from codeflash.languages import is_javascript
 from codeflash.models.models import CodeString, CodeStringsMarkdown
 
 if TYPE_CHECKING:
@@ -717,7 +718,7 @@ def detect_unused_helper_functions(
 
     """
     # Skip this analysis for non-Python languages since we use Python's ast module
-    if function_to_optimize.language in ("javascript", "typescript"):
+    if is_javascript():
         logger.debug("Skipping unused helper function detection for JavaScript/TypeScript")
         return []
 

codeflash/discovery/discover_unit_tests.py

@@ -28,6 +28,7 @@
     module_name_from_file_path,
 )
 from codeflash.code_utils.compat import SAFE_SYS_EXECUTABLE, codeflash_cache_db
+from codeflash.languages import is_javascript, is_python
 from codeflash.code_utils.shell_utils import get_cross_platform_subprocess_run_args
 from codeflash.models.models import CodePosition, FunctionCalledInTest, TestsInFile, TestType
 
@@ -659,10 +660,10 @@ def discover_unit_tests(
     language = _detect_language_from_functions(file_to_funcs_to_optimize)
 
     # Route to language-specific test discovery for non-Python languages
-    if language and language != "python":
+    if not is_python():
         # For JavaScript/TypeScript, tests_project_rootdir should be tests_root itself
         # The Jest helper will be configured to NOT include "tests." prefix to match
-        if language in ("javascript", "typescript"):
+        if is_javascript():
             cfg.tests_project_rootdir = cfg.tests_root
         return discover_tests_for_language(cfg, language, file_to_funcs_to_optimize)
 

codeflash/languages/__init__.py

@@ -28,6 +28,15 @@
     TestInfo,
     TestResult,
 )
+from codeflash.languages.current import (
+    current_language,
+    current_language_support,
+    is_javascript,
+    is_python,
+    is_typescript,
+    reset_current_language,
+    set_current_language,
+)
 from codeflash.languages.registry import (
     detect_project_language,
     get_language_support,
@@ -43,18 +52,26 @@
 
 __all__ = [
     # Base types
+    "CodeContext",
+    "FunctionInfo",
+    "HelperFunction",
     "Language",
     "LanguageSupport",
-    "FunctionInfo",
     "ParentInfo",
-    "CodeContext",
-    "HelperFunction",
-    "TestResult",
     "TestInfo",
+    "TestResult",
+    # Current language singleton
+    "current_language",
+    "current_language_support",
+    "is_javascript",
+    "is_python",
+    "is_typescript",
+    "reset_current_language",
+    "set_current_language",
     # Registry functions
-    "get_language_support",
     "detect_project_language",
-    "register_language",
-    "get_supported_languages",
+    "get_language_support",
     "get_supported_extensions",
+    "get_supported_languages",
+    "register_language",
 ]

codeflash/languages/current.py

@@ -0,0 +1,113 @@
+"""Singleton for the current language being used in the codeflash session.
+
+This module provides a centralized way to access and set the current language
+throughout the codeflash codebase, eliminating scattered language checks and
+string comparisons.
+
+Usage:
+    from codeflash.languages import current_language, set_current_language, is_python
+
+    # Set the language at the start of a session
+    set_current_language(Language.PYTHON)
+    # or
+    set_current_language("javascript")
+
+    # Check the current language anywhere in the codebase
+    if is_python():
+        # Python-specific code
+        ...
+
+    # Get the current language
+    lang = current_language()
+
+    # Get language support for the current language
+    support = current_language_support()
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from codeflash.languages.base import Language
+
+if TYPE_CHECKING:
+    from codeflash.languages.base import LanguageSupport
+
+# Module-level singleton for the current language
+_current_language: Language = Language.PYTHON
+
+
+def current_language() -> Language:
+    """Get the current language being used in this codeflash session.
+
+    Returns:
+        The current Language enum value.
+    """
+    return _current_language
+
+
+def set_current_language(language: Language | str) -> None:
+    """Set the current language for this codeflash session.
+
+    This should be called once at the start of an optimization run,
+    typically after reading the project configuration.
+
+    Args:
+        language: Either a Language enum value or a string like "python", "javascript", "typescript".
+    """
+    global _current_language
+
+    if isinstance(language, str):
+        _current_language = Language(language)
+    else:
+        _current_language = language
+
+
+def reset_current_language() -> None:
+    """Reset the current language to the default (Python).
+
+    Useful for testing or when starting a new session.
+    """
+    global _current_language
+    _current_language = Language.PYTHON
+
+
+def is_python() -> bool:
+    """Check if the current language is Python.
+
+    Returns:
+        True if the current language is Python.
+    """
+    return _current_language == Language.PYTHON
+
+
+def is_javascript() -> bool:
+    """Check if the current language is JavaScript or TypeScript.
+
+    This returns True for both JavaScript and TypeScript since they are
+    typically treated the same way in the optimization pipeline.
+
+    Returns:
+        True if the current language is JavaScript or TypeScript.
+    """
+    return _current_language in (Language.JAVASCRIPT, Language.TYPESCRIPT)
+
+
+def is_typescript() -> bool:
+    """Check if the current language is TypeScript specifically.
+
+    Returns:
+        True if the current language is TypeScript.
+    """
+    return _current_language == Language.TYPESCRIPT
+
+
+def current_language_support() -> LanguageSupport:
+    """Get the LanguageSupport instance for the current language.
+
+    Returns:
+        The LanguageSupport instance for the current language.
+    """
+    from codeflash.languages.registry import get_language_support
+
+    return get_language_support(_current_language)