perf: backport libcst visitor dispatch cache from codeflash-python

Cache the visitor dispatch tables that libcst rebuilds on every
MatcherDecoratableTransformer/Visitor instantiation. The tables
depend only on the class, not the instance, so caching by type is
safe. Saves ~27ms per visitor instantiation (24x faster).

Also fix pre-existing ruff F821 in cli.py (missing exit_with_message
import in process_pyproject_config).
This commit is contained in:
Kevin Turcios 2026-04-09 23:46:33 -05:00
parent 61053be9ce
commit b533f50bdc
13 changed files with 78 additions and 1 deletions

View file

@ -4,6 +4,7 @@ from typing import TYPE_CHECKING, Optional, Union
import libcst as cst
import codeflash.code_utils._libcst_cache # noqa: F401
from codeflash.code_utils.formatter import sort_imports
if TYPE_CHECKING:

View file

@ -82,7 +82,7 @@ def process_and_validate_cmd_args(args: Namespace) -> Namespace:
def process_pyproject_config(args: Namespace) -> Namespace:
from codeflash.code_utils import env_utils
from codeflash.code_utils.code_utils import normalize_ignore_paths
from codeflash.code_utils.code_utils import exit_with_message, normalize_ignore_paths
from codeflash.code_utils.config_parser import parse_config_file
from codeflash.languages.test_framework import set_current_test_framework
from codeflash.lsp.helpers import is_LSP_enabled

View file

@ -0,0 +1,64 @@
"""Cache libcst visitor dispatch table construction.
libcst's ``MatcherDecoratableTransformer`` and
``MatcherDecoratableVisitor`` rebuild visitor dispatch tables on
every instantiation by iterating ``dir(self)`` (~600 attributes)
and calling ``getattr`` + ``inspect.ismethod`` on each. The
results depend only on the class, not the instance, so caching
by ``type(obj)`` is safe.
Import this module before any libcst visitors are instantiated
to install the cache.
"""
from __future__ import annotations
from typing import Any
import libcst.matchers._visitors as _mv
_visit_cache: dict[type, Any] = {}
_leave_cache: dict[type, Any] = {}
_matchers_cache: dict[type, Any] = {}
_original_visit = _mv._gather_constructed_visit_funcs # noqa: SLF001
_original_leave = _mv._gather_constructed_leave_funcs # noqa: SLF001
_original_matchers = _mv._gather_matchers # noqa: SLF001
def _cached_visit(obj: object) -> Any:
"""Return cached visit-function dispatch table for the object's class."""
cls = type(obj)
try:
return _visit_cache[cls]
except KeyError:
result = _original_visit(obj)
_visit_cache[cls] = result
return result
def _cached_leave(obj: object) -> Any:
"""Return cached leave-function dispatch table for the object's class."""
cls = type(obj)
try:
return _leave_cache[cls]
except KeyError:
result = _original_leave(obj)
_leave_cache[cls] = result
return result
def _cached_matchers(obj: object) -> Any:
"""Return cached matcher dispatch table for the object's class."""
cls = type(obj)
try:
return dict(_matchers_cache[cls])
except KeyError:
result = _original_matchers(obj)
_matchers_cache[cls] = result
return dict(result)
_mv._gather_constructed_visit_funcs = _cached_visit # noqa: SLF001
_mv._gather_constructed_leave_funcs = _cached_leave # noqa: SLF001
_mv._gather_matchers = _cached_matchers # noqa: SLF001

View file

@ -7,6 +7,7 @@ from typing import TYPE_CHECKING
import libcst as cst
import codeflash.code_utils._libcst_cache # noqa: F401
from codeflash.cli_cmds.console import logger
from codeflash.code_utils.code_utils import get_run_tmp_file, module_name_from_file_path
from codeflash.code_utils.formatter import sort_imports

View file

@ -20,6 +20,7 @@ from rich.syntax import Syntax
from rich.text import Text
from rich.tree import Tree
import codeflash.code_utils._libcst_cache # noqa: F401
from codeflash.api.aiservice import AiServiceClient, AIServiceRefinerRequest, LocalAiServiceClient
from codeflash.api.cfapi import add_code_context_hash, create_staging, get_cfapi_base_urls, mark_optimization_success
from codeflash.benchmarking.utils import process_benchmark_data

View file

@ -11,6 +11,7 @@ from typing import TYPE_CHECKING, Any
import libcst as cst
import codeflash.code_utils._libcst_cache # noqa: F401
from codeflash.cli_cmds.console import logger
from codeflash.code_utils.code_utils import encoded_tokens_len, get_qualified_name, path_belongs_to_site_packages
from codeflash.code_utils.config_consts import (

View file

@ -9,6 +9,7 @@ from typing import TYPE_CHECKING, Optional, Union
import libcst as cst
import codeflash.code_utils._libcst_cache # noqa: F401
from codeflash.cli_cmds.console import logger
from codeflash.languages import current_language
from codeflash.languages.base import Language

View file

@ -11,6 +11,7 @@ from libcst.codemod import CodemodContext
from libcst.codemod.visitors import AddImportsVisitor, GatherImportsVisitor, RemoveImportsVisitor
from libcst.helpers import calculate_module_and_package
import codeflash.code_utils._libcst_cache # noqa: F401
from codeflash.cli_cmds.console import logger
from codeflash.code_utils.config_consts import MAX_CONTEXT_LEN_REVIEW
from codeflash.languages.base import Language

View file

@ -9,6 +9,7 @@ from typing import TYPE_CHECKING, TypeVar
import libcst as cst
from libcst.metadata import PositionProvider
import codeflash.code_utils._libcst_cache # noqa: F401
from codeflash.cli_cmds.console import logger
from codeflash.code_utils.config_parser import find_conftest_files
from codeflash.code_utils.formatter import sort_imports

View file

@ -10,6 +10,7 @@ import libcst as cst
from libcst import MetadataWrapper
from libcst.metadata import PositionProvider
import codeflash.code_utils._libcst_cache # noqa: F401
from codeflash.cli_cmds.console import logger
from codeflash.code_utils.time_utils import format_perf, format_time
from codeflash.models.models import GeneratedTests, GeneratedTestsList

View file

@ -9,6 +9,7 @@ from typing import TYPE_CHECKING, Union
import libcst as cst
import codeflash.code_utils._libcst_cache # noqa: F401
from codeflash.code_utils.code_utils import get_run_tmp_file
from codeflash.code_utils.formatter import sort_imports

View file

@ -9,6 +9,7 @@ from typing import TYPE_CHECKING, Any
import libcst as cst
import codeflash.code_utils._libcst_cache # noqa: F401
from codeflash.discovery.functions_to_optimize import FunctionToOptimize
from codeflash.languages.base import (
CodeContext,

View file

@ -19,6 +19,9 @@ from codeflash.models.test_type import TestType
if TYPE_CHECKING:
from collections.abc import Iterator
import libcst as cst
from rich.tree import Tree
@dataclass(frozen=True)
class AIServiceRefinerRequest: