⚡️ Speed up method CancelScope.__enter__ by 15%

[codeflash-ai]

📄 15% (0.15x) speedup for CancelScope.__enter__ in src/anyio/_backends/_asyncio.py

⏱️ Runtime : 1.69 milliseconds → 1.46 milliseconds (best of 74 runs)

📝 Explanation and details

The optimized code achieves a 15% speedup through several key optimizations:

1. WeakKeyDictionary Access Pattern (12.9% vs 10.5% time)

• Replaced try/except KeyError with _task_states.get(host_task) in __enter__
• Direct dictionary .get() is faster than exception handling for the common case where tasks already exist

2. Reduced Attribute Lookups in Hot Loops

• In _deliver_cancellation, cached self._tasks, self._host_task, and origin._cancel_reason as local variables
• Eliminated repeated attribute access during loop iterations (6% faster per-hit on task iteration)

3. Fast-path for math.inf Deadline Check

• Moved deadline comparison logic inline in __enter__ instead of always calling _timeout()
• Added early return in _timeout() for math.inf case
• Reduces function call overhead for the most common case (no timeout)

4. Optimized Attribute Access with getattr()

• Replaced task._must_cancel and task._fut_waiter with getattr() calls to avoid type ignore comments
• Slightly cleaner attribute access pattern

5. Branch Prediction Improvements

• Restructured conditional logic to avoid complex expressions in loop conditions
• Set should_retry = True only when cancellation is actually attempted

The optimizations are most effective for high-frequency context manager usage and nested scope hierarchies, as shown by the line profiler results where __enter__ time reduced from 11.6ms to 8.6ms total. The WeakKeyDictionary optimization alone accounts for most of the performance gain in typical async workloads.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

import asyncio
import math
import sys
from weakref import WeakKeyDictionary

# imports
import pytest
from anyio._backends._asyncio import CancelScope


# Minimal mocks and helpers to allow CancelScope to work in tests
class DummyTask:
    """A dummy asyncio.Task replacement for unit tests."""
    def __init__(self):
        self._cancelled = False
        self._must_cancel = False
        self._fut_waiter = DummyFuture()
        self._uncancel_count = 0

    def cancel(self, reason=None):
        self._cancelled = True
        self._cancel_reason = reason

    def uncancel(self):
        self._uncancel_count += 1

    def get_coro(self):
        return self

# Dummy Future for waiter
class DummyFuture:
    def __init__(self):
        self._done = False
    def done(self):
        return self._done

# Dummy current_task and get_running_loop
_current_task = None

def set_current_task(task):
    global _current_task
    _current_task = task

class DummyLoop:
    def __init__(self):
        self._time = 0
        self._scheduled = []
    def time(self):
        return self._time
    def call_at(self, deadline, callback):
        self._scheduled.append((deadline, callback))
        return DummyHandle()
    def call_soon(self, callback, *args):
        self._scheduled.append((self._time, lambda: callback(*args)))
        return DummyHandle()

class DummyHandle:
    def __init__(self):
        self._cancelled = False
    def cancel(self):
        self._cancelled = True

_dummy_loop = DummyLoop()

_task_states = WeakKeyDictionary()

# -------------------- UNIT TESTS --------------------

# Basic Test Cases













def test_enter_with_no_current_task():
    """Test __enter__ raises if current_task() returns None."""
    set_current_task(None)
    scope = CancelScope()
    with pytest.raises(TypeError):
        scope.__enter__()











#------------------------------------------------
import asyncio
import math
import sys
from types import TracebackType
# WeakKeyDictionary is used for _task_states
from weakref import WeakKeyDictionary

# Patch current_task for our tests
import anyio._backends._asyncio as tested_module
# imports
import pytest
from anyio._backends._asyncio import CancelScope

# --- Minimal stubs and helpers to allow CancelScope to be tested ---

# Simulate TaskState and _task_states as in the source code
class TaskState:
    def __init__(self, parent, cancel_scope):
        self.parent = parent
        self.cancel_scope = cancel_scope

_task_states = WeakKeyDictionary()

# Simulate current_task for synchronous tests
def fake_current_task():
    # Just return a dummy object with required attributes
    return DummyTask()

tested_module.current_task = fake_current_task

# Dummy asyncio.Task implementation for testing
class DummyTask:
    def __init__(self):
        self._must_cancel = False
        self._fut_waiter = asyncio.Future()
        self.cancelled = False
        self.cancel_reason = None
        self.uncancel_calls = 0

    def get_coro(self):
        return lambda: None  # dummy coroutine

    def cancel(self, reason=None):
        self.cancelled = True
        self.cancel_reason = reason

    def uncancel(self):
        self.uncancel_calls += 1

# Patch get_running_loop for our tests
class DummyLoop:
    def __init__(self):
        self._time = 0.0
        self.scheduled = []

    def time(self):
        return self._time

    def call_at(self, deadline, callback):
        self.scheduled.append((deadline, callback))
        return DummyHandle()

    def call_soon(self, callback, *args):
        self.scheduled.append((self._time, lambda: callback(*args)))
        return DummyHandle()

class DummyHandle:
    def cancel(self):
        pass

tested_module.get_running_loop = lambda: DummyLoop()

# Patch getcoroutinestate for our tests
tested_module.getcoroutinestate = lambda coro: tested_module.CORO_RUNNING

# Patch CORO_RUNNING and CORO_SUSPENDED
tested_module.CORO_RUNNING = 1
tested_module.CORO_SUSPENDED = 2

# Patch CancelledError and iterate_exceptions for __exit__ tests
class DummyCancelledError(Exception):
    pass

tested_module.CancelledError = DummyCancelledError
tested_module.iterate_exceptions = lambda exc: [exc]

def is_anyio_cancellation(exc):
    # For testing, treat all DummyCancelledError as anyio cancellation
    return isinstance(exc, DummyCancelledError)
tested_module.is_anyio_cancellation = is_anyio_cancellation

# --- Import the CancelScope to test ---
CancelScope = tested_module.CancelScope

# --- Unit tests for __enter__ ---

# 1. Basic Test Cases











def test_enter_handles_attribute_error_in_task_started():
    """
    Edge: If getcoroutinestate raises AttributeError, exception is raised.
    """
    scope = CancelScope()
    dummy_task = fake_current_task()
    def bad_get_coro():
        raise AttributeError
    dummy_task.get_coro = bad_get_coro
    tested_module.current_task = lambda: dummy_task
    with pytest.raises(Exception):
        scope.__enter__()
    tested_module.current_task = fake_current_task  # restore






#------------------------------------------------
from anyio._backends._asyncio import CancelScope
import pytest

def test_CancelScope___enter__():
    with pytest.raises(RuntimeError, match='no\\ running\\ event\\ loop'):
        CancelScope.__enter__(CancelScope(shield=False))

To edit these changes git checkout codeflash/optimize-CancelScope.__enter__-mhfo6i1u and push.