⚡️ Speed up function _should_use_raw_project_class_context by 51% in PR #1838 (codeflash/optimize-pr1660-2026-03-16T19.13.06)

[codeflash-ai]

⚡️ This pull request contains optimizations for PR #1838

If you approve this dependent PR, these changes will be merged into the original PR branch codeflash/optimize-pr1660-2026-03-16T19.13.06.

This PR will be automatically closed if the original PR is merged.

---

📄 51% (0.51x) speedup for _should_use_raw_project_class_context in codeflash/languages/python/context/code_context_extractor.py

⏱️ Runtime : 1.43 milliseconds → 948 microseconds (best of 163 runs)

📝 Explanation and details

The optimization replaced any() generator expressions with explicit early-return for-loops in four helper functions (_is_namedtuple_class, _class_has_explicit_init, _has_descriptor_like_class_fields, and _has_non_property_method_decorator), eliminating the overhead of building generator objects and calling the any() builtin. Line profiler data shows _class_has_explicit_init dropped from 1.85 ms to 0.96 ms (48% faster), and _is_namedtuple_class improved from 97 µs to 53 µs (46% faster), because the optimized code avoids allocating iterator state and returns immediately upon finding a match instead of completing the generator. The 51% overall runtime improvement (1.43 ms → 948 µs) comes from these cumulative reductions in per-call overhead across thousands of invocations during AST traversal. Test suite confirms no behavioral changes across all edge cases including dataclasses, decorators, and size-limit boundaries.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 1056 Passed
⏪ Replay Tests	🔘 None Found
🔎 Concolic Coverage Tests	🔘 None Found
📊 Tests Coverage	100.0%

🌀 Click to see Generated Regression Tests

import ast

# imports
import pytest  # used for our unit tests
# import the function under test and constants from the real module
from codeflash.languages.python.context.code_context_extractor import (
    MAX_RAW_PROJECT_CLASS_BODY_ITEMS, MAX_RAW_PROJECT_CLASS_LINES,
    _should_use_raw_project_class_context)

# helper to parse source and return the first ClassDef node (raises if not present)
def _get_class_node(src: str) -> ast.ClassDef:
    """
    Parse Python source and return the first top-level ClassDef node.

    We rely on the real ast.parse to produce lineno and end_lineno attributes,
    which _should_use_raw_project_class_context asserts exist.
    """
    module = ast.parse(src)
    for node in module.body:
        if isinstance(node, ast.ClassDef):
            return node
    raise ValueError("No class definition in source")

def test_small_class_with_explicit_init_returns_true():
    # A small class containing an explicit __init__ should be considered "raw project" (True).
    src = """
class A:
    def __init__(self):
        self.x = 1
"""
    cls = _get_class_node(src)
    # No import aliases needed for this simple case
    assert _should_use_raw_project_class_context(cls, {}) is True # 3.37μs -> 2.08μs (61.6% faster)

def test_small_class_without_init_and_no_features_returns_false():
    # A minimal small class with no init, no decorators, no descriptor fields, and no special methods -> False
    src = """
class B:
    pass
"""
    cls = _get_class_node(src)
    assert _should_use_raw_project_class_context(cls, {}) is False # 4.85μs -> 3.15μs (54.1% faster)

def test_namedtuple_base_detected_without_import_aliases():
    # If a class inherits from NamedTuple by direct name, it should be detected as namedtuple -> True
    src = """
class C(NamedTuple):
    x: int
"""
    cls = _get_class_node(src)
    assert _should_use_raw_project_class_context(cls, {}) is True # 4.59μs -> 2.81μs (63.0% faster)

def test_namedtuple_detected_via_import_aliases():
    # If the base is an alias (e.g., NT) and import_aliases map resolves it to a NamedTuple, it should be detected
    src = """
class D(NT):
    x: int
"""
    cls = _get_class_node(src)
    # Simulate "from typing import NamedTuple as NT" resolution
    import_aliases = {"NT": "typing.NamedTuple"}
    assert _should_use_raw_project_class_context(cls, import_aliases) is True # 5.37μs -> 3.56μs (51.0% faster)

def test_dataclass_decorator_with_kwargs_is_detected():
    # A class decorated with @dataclass(... ) should be detected regardless of keyword values
    src = """
@dataclass(init=False, kw_only=True)
class E:
    x: int
"""
    cls = _get_class_node(src)
    assert _should_use_raw_project_class_context(cls, {}) is True # 6.71μs -> 5.74μs (17.0% faster)

def test_any_class_decorator_triggers_true():
    # Any decorator on the class (even unrelated to dataclass) should make the function return True
    src = """
@some_decorator
class F:
    pass
"""
    cls = _get_class_node(src)
    assert _should_use_raw_project_class_context(cls, {}) is True # 5.71μs -> 4.51μs (26.7% faster)

def test_descriptor_like_class_field_triggers_true():
    # A class attribute assigned from a Call is considered "descriptor-like" and should trigger True
    src = """
class G:
    descriptor = Descriptor()
"""
    cls = _get_class_node(src)
    assert _should_use_raw_project_class_context(cls, {}) is True # 4.92μs -> 3.11μs (58.4% faster)

def test_non_property_method_decorator_triggers_true():
    # A method with a decorator that's not @property (e.g., @staticmethod) should cause True
    src = """
class H:
    @staticmethod
    def util():
        return 42
"""
    cls = _get_class_node(src)
    assert _should_use_raw_project_class_context(cls, {}) is True # 6.43μs -> 5.10μs (26.1% faster)

def test_only_property_decorated_method_does_not_trigger_true():
    # A single @property method should NOT count as a non-property decorated method; with no other triggers -> False
    src = """
class I:
    @property
    def value(self):
        return 1
"""
    cls = _get_class_node(src)
    assert _should_use_raw_project_class_context(cls, {}) is False # 5.27μs -> 3.67μs (43.7% faster)

def test_dataclass_decorator_with_non_boolean_kwargs_still_detects_dataclass():
    # Even if dataclass decorator has non-boolean keyword values, the presence of @dataclass should be detected
    src = """
@dataclass(init='weird', kw_only=0)
class J:
    a: int
"""
    cls = _get_class_node(src)
    # The function should identify the dataclass decorator presence and return True
    assert _should_use_raw_project_class_context(cls, {}) is True # 6.46μs -> 5.74μs (12.6% faster)

def test_decorator_call_expression_affects_start_line_and_counting():
    # Ensure class start line is taken from decorators (which could make class_line_count larger).
    # Construct a class with many blank lines after decorator so that line counting would be affected,
    # but still remain under MAX_RAW_PROJECT_CLASS_LINES to keep other behavior deterministic.
    decorator_lines = "\n" * 2  # extra blank lines before class start (decorators will push start_line up)
    src = f"""
{decorator_lines}
@some_decorator
class K:
    pass
"""
    cls = _get_class_node(src)
    # Should be True because decorator_list exists (regardless of line counting nuance)
    assert _should_use_raw_project_class_context(cls, {}) is True # 5.48μs -> 4.53μs (21.0% faster)

def test_large_class_with_many_simple_assignments_exceeds_body_item_limit_and_returns_false():
    # Create a class with many Assign nodes (but none are Calls), exceeding the MAX_RAW_PROJECT_CLASS_BODY_ITEMS limit.
    # It should be considered not "small" and, since no other triggers are present, return False.
    many = MAX_RAW_PROJECT_CLASS_BODY_ITEMS + 50  # comfortably exceed the threshold
    # Build class body with many simple assignments (values are constants, not calls)
    assigns = "\n".join(f"    a_{i} = {i}" for i in range(many))
    src = f"""
class Large:
{assigns}
"""
    cls = _get_class_node(src)
    # Sanity check: the AST body length should exceed the configured small-item threshold
    assert len(cls.body) > MAX_RAW_PROJECT_CLASS_BODY_ITEMS # 18.1μs -> 15.6μs (16.2% faster)
    assert _should_use_raw_project_class_context(cls, {}) is False

def test_many_iterations_of_function_with_small_class_are_deterministic_and_fast():
    # Run the function many times (1000 iterations) on the same small class to ensure deterministic results
    src = """
class LoopTest:
    def __init__(self):
        self.x = 1
"""
    cls = _get_class_node(src)
    # Run in a loop up to 1000 times, ensuring result is always True (no side effects)
    for _ in range(1000):
        assert _should_use_raw_project_class_context(cls, {}) is True # 1.05ms -> 638μs (63.8% faster)

import ast

# imports
import pytest
from codeflash.languages.python.context.code_context_extractor import \
    _should_use_raw_project_class_context

# Test fixtures for creating AST nodes
@pytest.fixture
def empty_import_aliases():
    """Fixture providing an empty import aliases dictionary."""
    return {}

def parse_class_code(code: str) -> ast.ClassDef:
    """Helper function to parse Python code and extract the first class definition."""
    tree = ast.parse(code)
    for node in ast.walk(tree):
        if isinstance(node, ast.ClassDef):
            return node
    raise ValueError("No class definition found in code")

def test_small_class_with_explicit_init_returns_true(empty_import_aliases):
    """Test that a small class with explicit __init__ returns True."""
    code = """
class MyClass:
    def __init__(self):
        self.x = 1
"""
    class_node = parse_class_code(code)
    result = _should_use_raw_project_class_context(class_node, empty_import_aliases) # 2.88μs -> 1.99μs (44.8% faster)
    assert result is True

def test_small_namedtuple_class_returns_true(empty_import_aliases):
    """Test that a NamedTuple class returns True."""
    code = """
from typing import NamedTuple
class MyTuple(NamedTuple):
    x: int
    y: str
"""
    class_node = parse_class_code(code)
    import_aliases = {"NamedTuple": "typing.NamedTuple"}
    result = _should_use_raw_project_class_context(class_node, import_aliases) # 4.42μs -> 2.85μs (55.3% faster)
    assert result is True

def test_dataclass_returns_true(empty_import_aliases):
    """Test that a dataclass decorated class returns True."""
    code = """
from dataclasses import dataclass
@dataclass
class MyDataClass:
    x: int
    y: str
"""
    class_node = parse_class_code(code)
    import_aliases = {"dataclass": "dataclasses.dataclass"}
    result = _should_use_raw_project_class_context(class_node, import_aliases) # 5.27μs -> 4.23μs (24.6% faster)
    assert result is True

def test_class_with_decorator_returns_true(empty_import_aliases):
    """Test that a class with any decorator returns True."""
    code = """
@some_decorator
class MyClass:
    x = 1
"""
    class_node = parse_class_code(code)
    result = _should_use_raw_project_class_context(class_node, empty_import_aliases) # 5.68μs -> 4.81μs (18.1% faster)
    assert result is True

def test_class_with_descriptor_fields_returns_true(empty_import_aliases):
    """Test that a class with descriptor-like fields returns True."""
    code = """
class MyClass:
    x = property(lambda self: self._x)
    y = descriptor_call()
"""
    class_node = parse_class_code(code)
    result = _should_use_raw_project_class_context(class_node, empty_import_aliases) # 5.02μs -> 3.30μs (52.3% faster)
    assert result is True

def test_class_with_method_decorator_returns_true(empty_import_aliases):
    """Test that a class with a method that has non-property decorators returns True."""
    code = """
class MyClass:
    @some_decorator
    def my_method(self):
        pass
"""
    class_node = parse_class_code(code)
    result = _should_use_raw_project_class_context(class_node, empty_import_aliases) # 6.41μs -> 5.10μs (25.8% faster)
    assert result is True

def test_simple_class_without_special_features_returns_false(empty_import_aliases):
    """Test that a simple class without special features returns False."""
    code = """
class MyClass:
    x = 1
    y = 2
"""
    class_node = parse_class_code(code)
    result = _should_use_raw_project_class_context(class_node, empty_import_aliases) # 5.06μs -> 3.75μs (35.0% faster)
    assert result is False

def test_class_with_only_property_methods_returns_false(empty_import_aliases):
    """Test that a class with only property methods returns False."""
    code = """
class MyClass:
    @property
    def x(self):
        return self._x
"""
    class_node = parse_class_code(code)
    import_aliases = {"property": "builtins.property"}
    result = _should_use_raw_project_class_context(class_node, import_aliases) # 5.19μs -> 3.71μs (40.0% faster)
    assert result is False

def test_large_class_with_explicit_init_returns_false(empty_import_aliases):
    """Test that a large class with explicit __init__ returns False."""
    # Create a class that exceeds MAX_RAW_PROJECT_CLASS_LINES (40 lines)
    lines = ["class MyClass:"]
    lines.append("    def __init__(self):")
    for i in range(50):
        lines.append(f"        self.x{i} = {i}")
    code = "\n".join(lines)
    class_node = parse_class_code(code)
    result = _should_use_raw_project_class_context(class_node, empty_import_aliases) # 4.01μs -> 2.83μs (41.3% faster)
    assert result is False

def test_class_with_many_body_items_and_init_returns_false(empty_import_aliases):
    """Test that a class with many body items (>8) and __init__ returns False."""
    # Create a class with more than MAX_RAW_PROJECT_CLASS_BODY_ITEMS (8)
    lines = ["class MyClass:"]
    lines.append("    def __init__(self): pass")
    for i in range(10):
        lines.append(f"    x{i} = {i}")
    code = "\n".join(lines)
    class_node = parse_class_code(code)
    result = _should_use_raw_project_class_context(class_node, empty_import_aliases) # 6.84μs -> 5.66μs (20.9% faster)
    assert result is False

def test_empty_class_returns_false(empty_import_aliases):
    """Test that an empty class returns False."""
    code = """
class MyClass:
    pass
"""
    class_node = parse_class_code(code)
    result = _should_use_raw_project_class_context(class_node, empty_import_aliases) # 4.43μs -> 2.77μs (59.6% faster)
    assert result is False

def test_class_with_empty_init_returns_true(empty_import_aliases):
    """Test that a class with an empty __init__ returns True (still has explicit init)."""
    code = """
class MyClass:
    def __init__(self):
        pass
"""
    class_node = parse_class_code(code)
    result = _should_use_raw_project_class_context(class_node, empty_import_aliases) # 2.81μs -> 1.85μs (51.3% faster)
    assert result is True

def test_class_with_async_init_returns_true(empty_import_aliases):
    """Test that a class with async __init__ is recognized (as FunctionDef/AsyncFunctionDef)."""
    code = """
class MyClass:
    async def __init__(self):
        pass
"""
    class_node = parse_class_code(code)
    result = _should_use_raw_project_class_context(class_node, empty_import_aliases) # 2.77μs -> 1.80μs (53.3% faster)
    assert result is True

def test_class_with_async_decorated_method_returns_true(empty_import_aliases):
    """Test that a class with an async method having decorators returns True."""
    code = """
class MyClass:
    @decorator
    async def my_async_method(self):
        pass
"""
    class_node = parse_class_code(code)
    result = _should_use_raw_project_class_context(class_node, empty_import_aliases) # 6.51μs -> 4.96μs (31.3% faster)
    assert result is True

def test_namedtuple_with_imported_alias(empty_import_aliases):
    """Test NamedTuple recognition when imported with an alias."""
    code = """
class MyTuple(NT):
    pass
"""
    class_node = parse_class_code(code)
    import_aliases = {"NT": "typing.NamedTuple"}
    result = _should_use_raw_project_class_context(class_node, import_aliases) # 5.28μs -> 3.45μs (53.2% faster)
    assert result is True

def test_dataclass_with_init_false_returns_true(empty_import_aliases):
    """Test that a dataclass with init=False still returns True."""
    code = """
@dataclass(init=False)
class MyClass:
    x: int
"""
    class_node = parse_class_code(code)
    import_aliases = {"dataclass": "dataclasses.dataclass"}
    result = _should_use_raw_project_class_context(class_node, import_aliases) # 6.05μs -> 5.29μs (14.4% faster)
    assert result is True

def test_dataclass_with_kw_only_returns_true(empty_import_aliases):
    """Test that a dataclass with kw_only=True returns True."""
    code = """
@dataclass(kw_only=True)
class MyClass:
    x: int
"""
    class_node = parse_class_code(code)
    import_aliases = {"dataclass": "dataclasses.dataclass"}
    result = _should_use_raw_project_class_context(class_node, import_aliases) # 5.92μs -> 5.14μs (15.2% faster)
    assert result is True

def test_class_with_decorator_on_class_and_method_returns_true(empty_import_aliases):
    """Test that decorator on class takes precedence and returns True."""
    code = """
@class_decorator
class MyClass:
    def simple_method(self):
        pass
"""
    class_node = parse_class_code(code)
    result = _should_use_raw_project_class_context(class_node, empty_import_aliases) # 5.47μs -> 4.62μs (18.4% faster)
    assert result is True

def test_class_with_property_setter_returns_false(empty_import_aliases):
    """Test that property.setter is recognized as property-like and returns False."""
    code = """
class MyClass:
    @property
    def x(self):
        return self._x
    
    @x.setter
    def x(self, value):
        self._x = value
"""
    class_node = parse_class_code(code)
    import_aliases = {"property": "builtins.property"}
    result = _should_use_raw_project_class_context(class_node, import_aliases) # 8.40μs -> 7.15μs (17.4% faster)
    assert result is False

def test_class_with_property_deleter_returns_false(empty_import_aliases):
    """Test that property.deleter is recognized as property-like and returns False."""
    code = """
class MyClass:
    @property
    def x(self):
        return self._x
    
    @x.deleter
    def x(self):
        del self._x
"""
    class_node = parse_class_code(code)
    import_aliases = {"property": "builtins.property"}
    result = _should_use_raw_project_class_context(class_node, import_aliases) # 8.02μs -> 6.66μs (20.4% faster)
    assert result is False

def test_class_at_exactly_max_lines_returns_true(empty_import_aliases):
    """Test that a class at exactly MAX_RAW_PROJECT_CLASS_LINES with __init__ returns True."""
    # MAX_RAW_PROJECT_CLASS_LINES = 40
    lines = ["class MyClass:"]
    lines.append("    def __init__(self):")
    lines.append("        pass")
    # Fill to exactly 40 lines total (including class def and init def)
    for i in range(36):
        lines.append(f"        # comment {i}")
    code = "\n".join(lines)
    class_node = parse_class_code(code)
    result = _should_use_raw_project_class_context(class_node, empty_import_aliases) # 2.81μs -> 1.76μs (59.7% faster)
    assert result is True

def test_class_at_exactly_max_body_items_returns_true(empty_import_aliases):
    """Test that a class at exactly MAX_RAW_PROJECT_CLASS_BODY_ITEMS with __init__ returns True."""
    # MAX_RAW_PROJECT_CLASS_BODY_ITEMS = 8
    lines = ["class MyClass:"]
    lines.append("    def __init__(self): pass")
    for i in range(7):
        lines.append(f"    x{i} = {i}")
    code = "\n".join(lines)
    class_node = parse_class_code(code)
    result = _should_use_raw_project_class_context(class_node, empty_import_aliases) # 2.81μs -> 1.84μs (52.3% faster)
    assert result is True

def test_class_exceeding_max_body_items_with_init_returns_false(empty_import_aliases):
    """Test that a class exceeding MAX_RAW_PROJECT_CLASS_BODY_ITEMS with __init__ returns False."""
    lines = ["class MyClass:"]
    lines.append("    def __init__(self): pass")
    for i in range(8):
        lines.append(f"    x{i} = {i}")
    code = "\n".join(lines)
    class_node = parse_class_code(code)
    result = _should_use_raw_project_class_context(class_node, empty_import_aliases) # 6.46μs -> 5.19μs (24.5% faster)
    assert result is False

def test_method_with_classmethod_decorator_returns_true(empty_import_aliases):
    """Test that a method with @classmethod returns True."""
    code = """
class MyClass:
    @classmethod
    def my_method(cls):
        pass
"""
    class_node = parse_class_code(code)
    result = _should_use_raw_project_class_context(class_node, empty_import_aliases) # 6.20μs -> 4.83μs (28.4% faster)
    assert result is True

def test_method_with_staticmethod_decorator_returns_true(empty_import_aliases):
    """Test that a method with @staticmethod returns True."""
    code = """
class MyClass:
    @staticmethod
    def my_method():
        pass
"""
    class_node = parse_class_code(code)
    result = _should_use_raw_project_class_context(class_node, empty_import_aliases) # 5.90μs -> 4.53μs (30.3% faster)
    assert result is True

def test_namedtuple_with_dotted_import_alias(empty_import_aliases):
    """Test NamedTuple recognition when fully qualified in alias."""
    code = """
class MyTuple(NamedTuple):
    x: int
"""
    class_node = parse_class_code(code)
    import_aliases = {"NamedTuple": "typing.NamedTuple"}
    result = _should_use_raw_project_class_context(class_node, import_aliases) # 4.08μs -> 2.69μs (51.9% faster)
    assert result is True

def test_dataclass_with_non_bool_decorator_arg(empty_import_aliases):
    """Test dataclass with non-bool arguments (should ignore and still return True)."""
    code = """
@dataclass(init=some_var)
class MyClass:
    x: int
"""
    class_node = parse_class_code(code)
    import_aliases = {"dataclass": "dataclasses.dataclass"}
    result = _should_use_raw_project_class_context(class_node, import_aliases) # 5.59μs -> 4.95μs (13.0% faster)
    assert result is True

def test_class_with_single_attribute_returns_false(empty_import_aliases):
    """Test that a class with a single simple attribute returns False."""
    code = """
class MyClass:
    x = 1
"""
    class_node = parse_class_code(code)
    result = _should_use_raw_project_class_context(class_node, empty_import_aliases) # 4.54μs -> 3.20μs (42.0% faster)
    assert result is False

def test_class_with_multiple_decorators_returns_true(empty_import_aliases):
    """Test that a class with multiple decorators returns True."""
    code = """
@decorator1
@decorator2
@decorator3
class MyClass:
    x = 1
"""
    class_node = parse_class_code(code)
    result = _should_use_raw_project_class_context(class_node, empty_import_aliases) # 6.49μs -> 5.70μs (13.9% faster)
    assert result is True

def test_simple_method_without_decorator_returns_false(empty_import_aliases):
    """Test that simple methods without decorators don't trigger True."""
    code = """
class MyClass:
    def regular_method(self):
        return 42
"""
    class_node = parse_class_code(code)
    result = _should_use_raw_project_class_context(class_node, empty_import_aliases) # 4.68μs -> 3.29μs (42.4% faster)
    assert result is False

def test_class_with_init_and_regular_method_returns_true(empty_import_aliases):
    """Test that class with __init__ returns True even with other methods."""
    code = """
class MyClass:
    def __init__(self):
        self.x = 1
    
    def other_method(self):
        return self.x
"""
    class_node = parse_class_code(code)
    result = _should_use_raw_project_class_context(class_node, empty_import_aliases) # 2.79μs -> 1.79μs (55.9% faster)
    assert result is True

def test_descriptor_with_call_expression(empty_import_aliases):
    """Test that descriptor-like fields with Call expressions trigger True."""
    code = """
class MyClass:
    x = some_descriptor()
    y = another_descriptor()
"""
    class_node = parse_class_code(code)
    result = _should_use_raw_project_class_context(class_node, empty_import_aliases) # 4.98μs -> 3.25μs (53.4% faster)
    assert result is True

def test_annotated_assignment_with_call(empty_import_aliases):
    """Test that annotated assignments with call expressions trigger True."""
    code = """
class MyClass:
    x: int = descriptor_call()
"""
    class_node = parse_class_code(code)
    result = _should_use_raw_project_class_context(class_node, empty_import_aliases) # 4.78μs -> 2.98μs (60.6% faster)
    assert result is True

def test_property_with_complex_name_structure(empty_import_aliases):
    """Test property recognition with complex name structures."""
    code = """
class MyClass:
    @module.property
    def x(self):
        return 1
"""
    class_node = parse_class_code(code)
    import_aliases = {"module.property": "builtins.property"}
    result = _should_use_raw_project_class_context(class_node, import_aliases) # 6.49μs -> 4.98μs (30.4% faster)
    # property should not trigger True when it's the only decorator
    assert result is False

def test_large_class_with_many_attributes(empty_import_aliases):
    """Test a large class with 100+ attributes but no __init__."""
    lines = ["class MyClass:"]
    for i in range(100):
        lines.append(f"    attr_{i} = {i}")
    code = "\n".join(lines)
    class_node = parse_class_code(code)
    result = _should_use_raw_project_class_context(class_node, empty_import_aliases) # 28.3μs -> 25.2μs (12.3% faster)
    assert result is False

def test_very_large_class_with_init_exceeding_limits(empty_import_aliases):
    """Test a very large class with __init__ that greatly exceeds size limits."""
    lines = ["class MyClass:"]
    lines.append("    def __init__(self):")
    for i in range(500):
        lines.append(f"        self.x_{i} = {i}")
    code = "\n".join(lines)
    class_node = parse_class_code(code)
    result = _should_use_raw_project_class_context(class_node, empty_import_aliases) # 4.52μs -> 3.26μs (38.8% faster)
    assert result is False

def test_class_with_many_methods_no_init(empty_import_aliases):
    """Test a class with many methods but no __init__."""
    lines = ["class MyClass:"]
    for i in range(50):
        lines.append(f"    def method_{i}(self):")
        lines.append("        pass")
    code = "\n".join(lines)
    class_node = parse_class_code(code)
    result = _should_use_raw_project_class_context(class_node, empty_import_aliases) # 17.8μs -> 15.8μs (13.2% faster)
    assert result is False

def test_class_with_many_decorated_methods(empty_import_aliases):
    """Test a class with many methods having decorators."""
    lines = ["class MyClass:"]
    for i in range(50):
        lines.append(f"    @decorator_{i}")
        lines.append(f"    def method_{i}(self):")
        lines.append("        pass")
    code = "\n".join(lines)
    class_node = parse_class_code(code)
    result = _should_use_raw_project_class_context(class_node, empty_import_aliases) # 11.1μs -> 9.12μs (21.5% faster)
    # Should return True because of the first decorated method found
    assert result is True

def test_class_with_many_descriptor_fields(empty_import_aliases):
    """Test a class with many descriptor-like fields."""
    lines = ["class MyClass:"]
    for i in range(100):
        lines.append(f"    field_{i} = descriptor_call_{i}()")
    code = "\n".join(lines)
    class_node = parse_class_code(code)
    result = _should_use_raw_project_class_context(class_node, empty_import_aliases) # 4.24μs -> 2.67μs (59.0% faster)
    assert result is True

def test_class_with_many_property_methods(empty_import_aliases):
    """Test a class with many property methods (should return False)."""
    lines = ["class MyClass:"]
    for i in range(100):
        lines.append(f"    @property")
        lines.append(f"    def prop_{i}(self):")
        lines.append(f"        return {i}")
    code = "\n".join(lines)
    class_node = parse_class_code(code)
    import_aliases = {"property": "builtins.property"}
    result = _should_use_raw_project_class_context(class_node, import_aliases) # 51.8μs -> 48.2μs (7.46% faster)
    assert result is False

def test_performance_with_deeply_nested_imports_dict(empty_import_aliases):
    """Test performance with a large import aliases dictionary."""
    lines = ["class MyClass:"]
    lines.append("    x = 1")
    code = "\n".join(lines)
    class_node = parse_class_code(code)
    
    # Create a large import aliases dictionary
    large_aliases = {f"import_{i}": f"module.submodule.import_{i}" for i in range(1000)}
    
    result = _should_use_raw_project_class_context(class_node, large_aliases) # 5.08μs -> 3.76μs (35.2% faster)
    assert result is False

def test_class_with_init_and_many_methods_under_limits(empty_import_aliases):
    """Test a class with __init__ and multiple methods staying under size limits."""
    lines = ["class MyClass:"]
    lines.append("    def __init__(self): pass")
    lines.append("    def method1(self): pass")
    lines.append("    def method2(self): pass")
    lines.append("    def method3(self): pass")
    lines.append("    def method4(self): pass")
    lines.append("    def method5(self): pass")
    lines.append("    def method6(self): pass")
    code = "\n".join(lines)
    class_node = parse_class_code(code)
    result = _should_use_raw_project_class_context(class_node, empty_import_aliases) # 2.92μs -> 1.80μs (61.6% faster)
    assert result is True

def test_namedtuple_with_many_fields(empty_import_aliases):
    """Test a NamedTuple with many fields."""
    lines = ["class MyTuple(NamedTuple):"]
    for i in range(500):
        lines.append(f"    field_{i}: int")
    code = "\n".join(lines)
    class_node = parse_class_code(code)
    import_aliases = {"NamedTuple": "typing.NamedTuple"}
    result = _should_use_raw_project_class_context(class_node, import_aliases) # 3.62μs -> 2.35μs (53.7% faster)
    assert result is True

def test_dataclass_with_many_fields(empty_import_aliases):
    """Test a dataclass with many fields."""
    lines = ["@dataclass"]
    lines.append("class MyDataClass:")
    for i in range(500):
        lines.append(f"    field_{i}: int = 0")
    code = "\n".join(lines)
    class_node = parse_class_code(code)
    import_aliases = {"dataclass": "dataclasses.dataclass"}
    result = _should_use_raw_project_class_context(class_node, import_aliases) # 4.89μs -> 4.12μs (18.8% faster)
    assert result is True

To edit these changes git checkout codeflash/optimize-pr1838-2026-03-16T19.35.10 and push.

[claude]

Claude finished @codeflash-ai[bot]'s task in 10s —— View job

---

PR Review Summary

• Triage: Assess PR scope — substantive code change, full review performed
• Lint & typecheck — fixed SIM110, committed & pushed
• Resolve stale threads — none found
• Code review
• Duplicate detection
• Test coverage
• Check/merge optimization PRs

---

Prek Checks

Fixed: ruff SIM110 flagged _is_namedtuple_class (line 777) because the optimization intentionally replaces any() with an explicit for-loop. Added # noqa: SIM110 to suppress it. Committed as style: auto-fix ruff SIM110 linting issue in _is_namedtuple_class and pushed.

mypy: No type errors.

Code Review

The optimization is sound. Four functions had any() generator expressions replaced with explicit early-return for-loops:

Function	Change
_is_namedtuple_class (line 776)	any(...) → for-loop
_class_has_explicit_init (line 817)	any(...) → for-loop
_has_descriptor_like_class_fields (line 953)	any(...) → for-loop
_has_non_property_method_decorator (line 940)	No change — already a for-loop in parent branch

No bugs, no correctness issues. The for-loop semantics are identical to any() for these patterns — both short-circuit on the first match. The performance improvement (51% overall, 48% for _class_has_explicit_init) is plausible given the elimination of generator allocation overhead across thousands of AST traversal calls.

Duplicate Detection

No duplicates detected. These helper functions are Python-AST-specific and only defined in codeflash/languages/python/context/code_context_extractor.py.

Test Coverage

1056 generated regression tests pass with 100% coverage per the PR description. No existing unit tests found (pre-existing condition). Coverage analysis skipped to avoid long-running test suite.

Open Optimization PRs

• PR ⚡️ Speed up function collect_existing_class_names by 188% in PR #1660 (unstructured-inference) #1838 (codeflash/optimize-pr1660-2026-03-16T19.13.06): CI not yet available — this is the parent PR that the current PR targets. Left open.
• PR ⚡️ Speed up function _parse_and_collect_imports by 12% in PR #1660 (unstructured-inference) #1837 (codeflash/optimize-pr1660-2026-03-16T18.33.59): CI not yet available. Left open.

---

| Branch: codeflash/optimize-pr1838-2026-03-16T19.35.10

[codeflash-ai]

⚡️Codeflash found 29% (0.29x) speedup for _is_namedtuple_class in codeflash/languages/python/context/code_context_extractor.py

⏱️ Runtime : 447 microseconds → 347 microseconds (best of 205 runs)

⚡️ This change will improve the performance of the following benchmarks:

Benchmark File :: Function	Original Runtime	Expected New Runtime	Speedup
tests.benchmarks.test_benchmark_code_extract_code_context::test_benchmark_extract	17.6 seconds	17.6 seconds	0.00%

🔻 This change will degrade the performance of the following benchmarks:

{benchmark_info_degraded}

📝 Explanation and details

The optimization inlines _expr_matches_name directly into _is_namedtuple_class, eliminating 2,150 function calls that originally consumed 92.6% of total execution time. The hot loop now computes string concatenations (".NamedTuple") once outside the loop instead of on every iteration, and restructures the resolved-alias check to short-circuit earlier when no alias exists. Line profiler confirms per-iteration cost dropped from ~3,786 ns to ~1,943 ns. No test cases regress in correctness, and the 28% runtime improvement directly benefits callers like _collect_synthetic_constructor_type_names and _build_synthetic_init_stub that check hundreds of class definitions during AST traversal.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 53 Passed
⏪ Replay Tests	✅ 1 Passed
🔎 Concolic Coverage Tests	🔘 None Found
📊 Tests Coverage	100.0%

🌀 Click to see Generated Regression Tests

import ast  # used to construct AST nodes that simulate class definitions and bases

import pytest  # used for our unit tests
# import the function to test from the provided module path
from codeflash.languages.python.context.code_context_extractor import \
    _is_namedtuple_class
from jedi.api.classes import Name

def test_namedtuple_with_simple_name_is_detected():
    # Create a ClassDef node whose base is a simple Name node "NamedTuple"
    class_node = ast.ClassDef(
        name="C",  # arbitrary class name
        bases=[ast.Name(id="NamedTuple")],  # base: NamedTuple
        keywords=[],  # no keywords
        body=[],  # empty body
        decorator_list=[],  # no decorators
    )
    # No import aliases are required for this direct match
    assert _is_namedtuple_class(class_node, {}) is True # 1.17μs -> 1.02μs (14.7% faster)

def test_namedtuple_with_qualified_name_is_detected():
    # Create a ClassDef node whose base is an attribute like "typing.NamedTuple"
    base_attr = ast.Attribute(value=ast.Name(id="typing"), attr="NamedTuple")
    class_node = ast.ClassDef(
        name="C",
        bases=[base_attr],  # base: typing.NamedTuple
        keywords=[],
        body=[],
        decorator_list=[],
    )
    # Should detect Qualified.NamedTuple because endswith(".NamedTuple")
    assert _is_namedtuple_class(class_node, {}) is True # 2.17μs -> 1.87μs (16.1% faster)

def test_namedtuple_when_base_is_call_is_detected():
    # Create a ClassDef node whose base is a Call node: NamedTuple(...)
    call_node = ast.Call(func=ast.Name(id="NamedTuple"), args=[], keywords=[])
    class_node = ast.ClassDef(
        name="C",
        bases=[call_node],  # base: NamedTuple(...) form
        keywords=[],
        body=[],
        decorator_list=[],
    )
    # The function should extract the func name from the Call and detect NamedTuple
    assert _is_namedtuple_class(class_node, {}) is True # 1.40μs -> 1.28μs (9.44% faster)

def test_namedtuple_detected_via_import_alias():
    # Simulate: from typing import NamedTuple as NT  --> base AST uses "NT"
    class_node = ast.ClassDef(
        name="C",
        bases=[ast.Name(id="NT")],  # base: NT (alias)
        keywords=[],
        body=[],
        decorator_list=[],
    )
    # import_aliases maps alias "NT" to the fully-qualified "typing.NamedTuple"
    import_aliases = {"NT": "typing.NamedTuple"}
    # Should resolve alias and detect NamedTuple
    assert _is_namedtuple_class(class_node, import_aliases) is True # 1.89μs -> 1.55μs (22.0% faster)

def test_non_namedtuple_base_returns_false():
    # Base is a different name that should not match
    class_node = ast.ClassDef(
        name="C",
        bases=[ast.Name(id="BaseClass")],  # base: BaseClass
        keywords=[],
        body=[],
        decorator_list=[],
    )
    # No alias mapping; should not detect NamedTuple
    assert _is_namedtuple_class(class_node, {}) is False # 1.59μs -> 1.28μs (24.2% faster)

def test_empty_bases_returns_false():
    # Class with no bases at all should not be treated as NamedTuple
    class_node = ast.ClassDef(
        name="C",
        bases=[],  # no bases provided
        keywords=[],
        body=[],
        decorator_list=[],
    )
    # Should be False because there are no bases to match
    assert _is_namedtuple_class(class_node, {}) is False # 421ns -> 450ns (6.44% slower)

def test_unknown_node_types_in_bases_are_ignored():
    # Some AST nodes do not have names; e.g., a Constant used incorrectly as a base
    class_node = ast.ClassDef(
        name="C",
        bases=[ast.Constant(value=42)],  # nonsensical base but valid AST node
        keywords=[],
        body=[],
        decorator_list=[],
    )
    # The function should ignore unknown node types and return False
    assert _is_namedtuple_class(class_node, {}) is False # 1.27μs -> 1.10μs (15.4% faster)

def test_alias_resolution_is_case_sensitive_and_mismatch_returns_false():
    # Alias maps to a lowercase "namedtuple" (e.g., collections.namedtuple) which should not match "NamedTuple"
    class_node = ast.ClassDef(
        name="C",
        bases=[ast.Name(id="NT")],  # alias used in base
        keywords=[],
        body=[],
        decorator_list=[],
    )
    import_aliases = {"NT": "collections.namedtuple"}  # not the same suffix, case differs
    # Should not match because resolved name does not end with ".NamedTuple"
    assert _is_namedtuple_class(class_node, import_aliases) is False # 1.94μs -> 1.57μs (23.6% faster)

def test_nested_attribute_like_pkg_sub_namedtuple_is_detected():
    # Create a nested attribute pkg.sub.NamedTuple
    nested_attr = ast.Attribute(
        value=ast.Attribute(value=ast.Name(id="pkg"), attr="sub"),
        attr="NamedTuple",
    )
    class_node = ast.ClassDef(
        name="C",
        bases=[nested_attr],  # base: pkg.sub.NamedTuple
        keywords=[],
        body=[],
        decorator_list=[],
    )
    # Should detect because final attribute name ends with ".NamedTuple"
    assert _is_namedtuple_class(class_node, {}) is True # 2.40μs -> 2.17μs (10.6% faster)

def test_call_of_attribute_with_import_alias_resolution():
    # Simulate base: mod.NT(...) where mod.NT is an alias to typing.NamedTuple
    call_node = ast.Call(func=ast.Attribute(value=ast.Name(id="mod"), attr="NT"), args=[], keywords=[])
    class_node = ast.ClassDef(
        name="C",
        bases=[call_node],  # base: mod.NT(...)
        keywords=[],
        body=[],
        decorator_list=[],
    )
    # Provide import_aliases that map the attribute name "mod.NT" to "typing.NamedTuple"
    import_aliases = {"mod.NT": "typing.NamedTuple"}
    # Should resolve through the attribute and detect NamedTuple
    assert _is_namedtuple_class(class_node, import_aliases) is True # 2.80μs -> 2.33μs (20.3% faster)

def test_passing_none_for_import_aliases_raises_attribute_error():
    # If import_aliases is None, the code attempts to call .get and should raise AttributeError
    class_node = ast.ClassDef(
        name="C",
        bases=[ast.Name(id="Other")],  # triggers the branch that will try to consult import_aliases.get
        keywords=[],
        body=[],
        decorator_list=[],
    )
    # Expect an AttributeError because NoneType has no .get method
    with pytest.raises(AttributeError):
        _is_namedtuple_class(class_node, None) # 4.77μs -> 4.17μs (14.4% faster)

def test_large_scale_mixed_classes_detection():
    # Construct a large list of ClassDef nodes (1000) with alternating matching and non-matching bases
    nodes = []
    import_aliases = {"AliasNT": "typing.NamedTuple"}  # provide one alias mapping used for some nodes
    total = 1000  # number of class nodes to generate (within the requested limit)
    # Populate nodes deterministically: even indices will match, odd will not
    for i in range(total):
        if i % 4 == 0:
            # direct NamedTuple name -> should match
            bases = [ast.Name(id="NamedTuple")]
        elif i % 4 == 1:
            # qualified attribute typing.NamedTuple -> should match
            bases = [ast.Attribute(value=ast.Name(id="typing"), attr="NamedTuple")]
        elif i % 4 == 2:
            # alias name that resolves to NamedTuple via import_aliases -> should match
            bases = [ast.Name(id="AliasNT")]
        else:
            # non-matching random name -> should not match
            bases = [ast.Name(id=f"Other{i}")]
        nodes.append(
            ast.ClassDef(
                name=f"C{i}",
                bases=bases,
                keywords=[],
                body=[],
                decorator_list=[],
            )
        )

    # Evaluate all nodes and count how many are detected as NamedTuple classes
    detected = 0
    for node in nodes:
        if _is_namedtuple_class(node, import_aliases):
            detected += 1

    # For our construction, 3 out of every 4 nodes should match
    expected = total * 3 // 4
    assert detected == expected  # ensure detection count matches expected value

import ast
from typing import NamedTuple

# imports
import pytest
from codeflash.languages.python.context.code_context_extractor import \
    _is_namedtuple_class

class TestIsNamedtupleClassBasic:
    """Basic tests for _is_namedtuple_class with typical use cases."""

    def test_namedtuple_class_with_direct_base(self):
        """Test that a class directly inheriting from NamedTuple is identified."""
        # Parse code defining a class that directly inherits from NamedTuple
        code = "class Point(NamedTuple): x: int; y: int"
        module = ast.parse(code)
        class_node = module.body[0]
        
        # Call the function with empty import_aliases (NamedTuple is a direct base)
        result = _is_namedtuple_class(class_node, {}) # 1.34μs -> 1.19μs (12.6% faster)
        
        # Should return True because the class has NamedTuple as a base
        assert result is True

    def test_namedtuple_class_with_imported_base(self):
        """Test that a class inheriting from an aliased NamedTuple is identified."""
        # Parse code defining a class that inherits from an aliased NamedTuple
        code = "class Point(MyTuple): x: int; y: int"
        module = ast.parse(code)
        class_node = module.body[0]
        
        # Provide an import_aliases mapping that resolves MyTuple to NamedTuple
        import_aliases = {"MyTuple": "typing.NamedTuple"}
        result = _is_namedtuple_class(class_node, import_aliases) # 2.31μs -> 1.80μs (28.4% faster)
        
        # Should return True because MyTuple resolves to NamedTuple
        assert result is True

    def test_non_namedtuple_class_with_simple_base(self):
        """Test that a class not inheriting from NamedTuple returns False."""
        # Parse code defining a regular class
        code = "class Point: x: int; y: int"
        module = ast.parse(code)
        class_node = module.body[0]
        
        # Call the function with empty import_aliases
        result = _is_namedtuple_class(class_node, {}) # 530ns -> 511ns (3.72% faster)
        
        # Should return False because the class doesn't inherit from NamedTuple
        assert result is False

    def test_class_with_non_namedtuple_base(self):
        """Test that a class inheriting from a non-NamedTuple base returns False."""
        # Parse code defining a class that inherits from a custom class
        code = "class Point(CustomBase): x: int; y: int"
        module = ast.parse(code)
        class_node = module.body[0]
        
        # Call the function with import_aliases that doesn't map to NamedTuple
        import_aliases = {"CustomBase": "my_module.CustomBase"}
        result = _is_namedtuple_class(class_node, import_aliases) # 2.31μs -> 1.79μs (29.1% faster)
        
        # Should return False because CustomBase doesn't resolve to NamedTuple
        assert result is False

    def test_namedtuple_with_fully_qualified_name(self):
        """Test that a class inheriting from typing.NamedTuple is identified."""
        # Parse code defining a class that inherits from fully qualified NamedTuple
        code = "class Point(typing.NamedTuple): x: int; y: int"
        module = ast.parse(code)
        class_node = module.body[0]
        
        # Call the function with empty import_aliases
        result = _is_namedtuple_class(class_node, {}) # 2.33μs -> 2.04μs (14.2% faster)
        
        # Should return True because typing.NamedTuple matches the pattern
        assert result is True

class TestIsNamedtupleClassEdgeCases:
    """Edge case tests for _is_namedtuple_class."""

    def test_class_with_no_bases(self):
        """Test that a class with no base classes returns False."""
        # Parse code defining a class with no explicit bases
        code = "class Empty: pass"
        module = ast.parse(code)
        class_node = module.body[0]
        
        # Call the function with empty import_aliases
        result = _is_namedtuple_class(class_node, {}) # 511ns -> 471ns (8.49% faster)
        
        # Should return False because the class has no bases
        assert result is False

    def test_class_with_empty_import_aliases(self):
        """Test that empty import_aliases dict works correctly."""
        # Parse code defining a class inheriting from NamedTuple
        code = "class Point(NamedTuple): pass"
        module = ast.parse(code)
        class_node = module.body[0]
        
        # Call the function with empty import_aliases
        result = _is_namedtuple_class(class_node, {}) # 1.17μs -> 972ns (20.6% faster)
        
        # Should return True because NamedTuple is recognized directly
        assert result is True

    def test_namedtuple_alias_single_suffix_match(self):
        """Test that an alias matching only the suffix is recognized."""
        # Parse code with an aliased base that matches the suffix
        code = "class Point(NT): pass"
        module = ast.parse(code)
        class_node = module.body[0]
        
        # Provide an import_aliases where NT resolves to a module with NamedTuple
        import_aliases = {"NT": "collections.namedtuple"}
        result = _is_namedtuple_class(class_node, import_aliases) # 2.20μs -> 1.64μs (34.1% faster)
        
        # Should return False because "namedtuple" doesn't end with "NamedTuple"
        assert result is False

    def test_namedtuple_from_collections_module(self):
        """Test that NamedTuple from collections module is recognized."""
        # Parse code with collections.NamedTuple
        code = "class Point(collections.NamedTuple): pass"
        module = ast.parse(code)
        class_node = module.body[0]
        
        # Call the function with empty import_aliases
        result = _is_namedtuple_class(class_node, {}) # 2.28μs -> 2.01μs (13.4% faster)
        
        # Should return True because it ends with NamedTuple
        assert result is True

    def test_class_with_multiple_bases_one_is_namedtuple(self):
        """Test that a class with multiple bases where one is NamedTuple is identified."""
        # Parse code defining a class with multiple bases
        code = "class Point(Mixin, NamedTuple): pass"
        module = ast.parse(code)
        class_node = module.body[0]
        
        # Call the function with empty import_aliases
        result = _is_namedtuple_class(class_node, {}) # 2.19μs -> 1.69μs (29.6% faster)
        
        # Should return True because one of the bases is NamedTuple
        assert result is True

    def test_class_with_multiple_bases_none_are_namedtuple(self):
        """Test that a class with multiple non-NamedTuple bases returns False."""
        # Parse code defining a class with multiple non-NamedTuple bases
        code = "class Point(Mixin1, Mixin2): pass"
        module = ast.parse(code)
        class_node = module.body[0]
        
        # Call the function with empty import_aliases
        result = _is_namedtuple_class(class_node, {}) # 2.21μs -> 1.84μs (20.1% faster)
        
        # Should return False because none of the bases are NamedTuple
        assert result is False

    def test_partial_name_match_should_fail(self):
        """Test that partial name matches don't incorrectly match."""
        # Parse code with a base named something containing NamedTuple
        code = "class Point(NotNamedTuple): pass"
        module = ast.parse(code)
        class_node = module.body[0]
        
        # Call the function with empty import_aliases
        result = _is_namedtuple_class(class_node, {}) # 1.74μs -> 1.44μs (20.8% faster)
        
        # Should return False because NotNamedTuple doesn't match "NamedTuple"
        assert result is False

    def test_case_sensitive_matching(self):
        """Test that matching is case-sensitive."""
        # Parse code with a base named namedtuple (lowercase)
        code = "class Point(namedtuple): pass"
        module = ast.parse(code)
        class_node = module.body[0]
        
        # Call the function with empty import_aliases
        result = _is_namedtuple_class(class_node, {}) # 1.77μs -> 1.46μs (21.3% faster)
        
        # Should return False because "namedtuple" != "NamedTuple"
        assert result is False

    def test_namedtuple_with_complex_qualified_name(self):
        """Test that deeply qualified NamedTuple names are recognized."""
        # Parse code with a complex qualified name
        code = "class Point(pkg.sub.mod.NamedTuple): pass"
        module = ast.parse(code)
        class_node = module.body[0]
        
        # Call the function with empty import_aliases
        result = _is_namedtuple_class(class_node, {}) # 2.88μs -> 2.72μs (5.85% faster)
        
        # Should return True because it ends with NamedTuple
        assert result is True

    def test_namedtuple_with_call_expression(self):
        """Test behavior when a base is a call expression to NamedTuple."""
        # Note: This is not valid Python syntax for class inheritance,
        # but the function should handle it gracefully
        code = "class Point(NamedTuple()): pass"
        module = ast.parse(code)
        class_node = module.body[0]
        
        # Call the function with empty import_aliases
        result = _is_namedtuple_class(class_node, {}) # 1.54μs -> 1.40μs (9.98% faster)
        
        # Should return True because _expr_matches_name handles Call nodes
        assert result is True

    def test_import_alias_override(self):
        """Test that import aliases can override base class names."""
        # Parse code with an aliased base
        code = "class Point(Tuple): pass"
        module = ast.parse(code)
        class_node = module.body[0]
        
        # Provide import_aliases where Tuple resolves to typing.NamedTuple
        import_aliases = {"Tuple": "typing.NamedTuple"}
        result = _is_namedtuple_class(class_node, import_aliases) # 2.19μs -> 1.69μs (29.6% faster)
        
        # Should return True because Tuple is aliased to NamedTuple
        assert result is True

    def test_import_alias_not_matching(self):
        """Test that import aliases not matching NamedTuple return False."""
        # Parse code with an aliased base
        code = "class Point(Tuple): pass"
        module = ast.parse(code)
        class_node = module.body[0]
        
        # Provide import_aliases where Tuple doesn't resolve to NamedTuple
        import_aliases = {"Tuple": "typing.Tuple"}
        result = _is_namedtuple_class(class_node, import_aliases) # 1.96μs -> 1.50μs (30.7% faster)
        
        # Should return False because Tuple is not aliased to NamedTuple
        assert result is False

class TestIsNamedtupleClassLargeScale:
    """Large-scale tests for _is_namedtuple_class."""

    def test_class_with_many_bases_one_is_namedtuple(self):
        """Test that a class with many bases where one is NamedTuple is identified."""
        # Generate code with many base classes, one of which is NamedTuple
        bases = ", ".join([f"Mixin{i}" for i in range(500)] + ["NamedTuple"])
        code = f"class Point({bases}): pass"
        module = ast.parse(code)
        class_node = module.body[0]
        
        # Call the function with empty import_aliases
        result = _is_namedtuple_class(class_node, {}) # 144μs -> 107μs (35.0% faster)
        
        # Should return True despite many bases
        assert result is True

    def test_class_with_many_non_namedtuple_bases(self):
        """Test that a class with many non-NamedTuple bases returns False."""
        # Generate code with many base classes, none of which are NamedTuple
        bases = ", ".join([f"Mixin{i}" for i in range(500)])
        code = f"class Point({bases}): pass"
        module = ast.parse(code)
        class_node = module.body[0]
        
        # Call the function with empty import_aliases
        result = _is_namedtuple_class(class_node, {}) # 143μs -> 104μs (37.2% faster)
        
        # Should return False
        assert result is False

    def test_large_import_aliases_dict_with_namedtuple_match(self):
        """Test that a large import_aliases dict with NamedTuple mapping works."""
        # Parse code with an aliased base
        code = "class Point(Alias500): pass"
        module = ast.parse(code)
        class_node = module.body[0]
        
        # Create a large import_aliases dict with many entries
        import_aliases = {f"Alias{i}": f"module{i}.Something" for i in range(1000)}
        # Set one entry to resolve to NamedTuple
        import_aliases["Alias500"] = "typing.NamedTuple"
        
        result = _is_namedtuple_class(class_node, import_aliases) # 2.50μs -> 1.90μs (31.0% faster)
        
        # Should return True despite large dict
        assert result is True

    def test_large_import_aliases_dict_without_namedtuple_match(self):
        """Test that a large import_aliases dict without NamedTuple mapping works."""
        # Parse code with an aliased base
        code = "class Point(Alias250): pass"
        module = ast.parse(code)
        class_node = module.body[0]
        
        # Create a large import_aliases dict with many entries
        import_aliases = {f"Alias{i}": f"module{i}.Something" for i in range(1000)}
        
        result = _is_namedtuple_class(class_node, import_aliases) # 2.29μs -> 1.76μs (30.1% faster)
        
        # Should return False
        assert result is False

    def test_many_qualified_namedtuple_variations(self):
        """Test multiple variations of qualified NamedTuple references."""
        # Create many test cases with different qualified NamedTuple names
        for depth in range(1, 20):
            qualifiers = ".".join([f"mod{i}" for i in range(depth)])
            code = f"class Point({qualifiers}.NamedTuple): pass"
            module = ast.parse(code)
            class_node = module.body[0]
            
            result = _is_namedtuple_class(class_node, {}) # 55.1μs -> 53.1μs (3.80% faster)
            
            # All should return True
            assert result is True, f"Failed for depth {depth}"

    def test_many_alias_lookups(self):
        """Test efficiency of many alias lookups."""
        # Parse code with multiple different bases
        bases = ", ".join([f"Type{i}" for i in range(100)])
        code = f"class Point({bases}): pass"
        module = ast.parse(code)
        class_node = module.body[0]
        
        # Create import_aliases for all but one (the last should be NamedTuple)
        import_aliases = {f"Type{i}": f"module{i}.Class{i}" for i in range(99)}
        import_aliases["Type99"] = "typing.NamedTuple"
        
        result = _is_namedtuple_class(class_node, import_aliases) # 43.1μs -> 31.5μs (36.8% faster)
        
        # Should return True by finding Type99
        assert result is True

⏪ Click to see Replay Tests

Test File::Test Function	Original ⏱️	Optimized ⏱️	Speedup
benchmarks/codeflash_replay_tests_gwziwff1/test_tests_benchmarks_test_benchmark_code_extract_code_context__replay_test_0.py::test_codeflash_languages_python_context_code_context_extractor__is_namedtuple_class_test_benchmark_extract	2.50μs	2.17μs	15.2%✅

To test or edit this optimization locally git merge codeflash/optimize-pr1839-2026-03-16T19.49.25

Suggested change

	for base in class_node.bases: # noqa: SIM110
	if _expr_matches_name(base, import_aliases, "NamedTuple"):
	return True
	dot_suffix = ".NamedTuple"
	for base in class_node.bases: # noqa: SIM110
	expr_name = _get_expr_name(base)
	if expr_name is None:
	continue
	if expr_name == "NamedTuple" or expr_name.endswith(dot_suffix):
	return True
	resolved_name = import_aliases.get(expr_name)
	if resolved_name is not None and (resolved_name == "NamedTuple" or resolved_name.endswith(dot_suffix)):
	return True

[codeflash-ai]

⚡️ Codeflash found optimizations for this PR

📄 11% (0.11x) speedup for _class_has_explicit_init in codeflash/languages/python/context/code_context_extractor.py

⏱️ Runtime : 494 microseconds → 444 microseconds (best of 148 runs)

A dependent PR with the suggested changes has been created. Please review:

• ⚡️ Speed up function _class_has_explicit_init by 11% in PR #1839 (codeflash/optimize-pr1838-2026-03-16T19.35.10) #1841

If you approve, it will be merged into this PR (branch codeflash/optimize-pr1838-2026-03-16T19.35.10).