⚡️ Speed up function pytest_sessionfinish by 125%

[codeflash-ai]

📄 125% (1.25x) speedup for pytest_sessionfinish in src/_pytest/tmpdir.py

⏱️ Runtime : 38.6 milliseconds → 17.2 milliseconds (best of 87 runs)

📝 Explanation and details

The optimization replaces the expensive left_dir.resolve().exists() check with a more efficient left_dir.stat() call wrapped in a try-except block to detect dead symlinks.

Key optimization in cleanup_dead_symlinks:

• Original approach: if not left_dir.resolve().exists() - This calls .resolve() which follows the symlink to its target, then .exists() to check if that target exists. For dead symlinks, .resolve() can be expensive as it tries to traverse non-existent paths.
• Optimized approach: try: left_dir.stat() except FileNotFoundError: - This directly attempts to stat the symlink target. If the target doesn't exist, it raises FileNotFoundError, which we catch to identify dead symlinks.

Why this is faster:
The line profiler shows the critical bottleneck was line if not left_dir.resolve().exists() taking 75.7% of total time (76.4ms) in the original version. The optimized version reduces this to just 16.6% (4.85ms) for the left_dir.stat() call, delivering a 3.7x speedup in the core symlink detection logic.

Performance impact by test case:

• Basic cases with few symlinks show minimal improvement (0-6% faster)
• Edge cases with dead symlinks show significant gains (120-175% faster)
• Large-scale tests with many symlinks show dramatic improvements (162-225% faster)

The optimization is particularly effective for workloads with many dead symlinks, as evidenced by the test_large_number_of_files_and_symlinks case improving from 17.5ms to 5.39ms (225% faster). Since this function is called during pytest session cleanup, the improvement reduces test suite teardown time, especially in environments with many temporary symlinks.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

# imports
from _pytest.tmpdir import pytest_sessionfinish


# --- Helpers for test doubles ---


class DummyConfig:
    def __init__(self, tmp_path_factory):
        self._tmp_path_factory = tmp_path_factory


class DummySession:
    def __init__(self, tmp_path_factory):
        self.config = DummyConfig(tmp_path_factory)


class DummyTmpPathFactory:
    def __init__(self, basetemp, retention_policy="failed", given_basetemp=None):
        self._basetemp = basetemp
        self._retention_policy = retention_policy
        self._given_basetemp = given_basetemp


# --- Unit Tests ---

# BASIC TEST CASES


def test_no_basetemp_returns(tmp_path):
    """If basetemp is None, function should return and do nothing."""
    factory = DummyTmpPathFactory(basetemp=None)
    session = DummySession(factory)
    pytest_sessionfinish(session, exitstatus=0)  # 844ns -> 818ns (3.18% faster)
    # Nothing to assert, but should not error


def test_policy_failed_exitstatus_0_removes_dir(tmp_path):
    """If exitstatus==0, policy=='failed', and no given_basetemp, remove the directory."""
    # Make a temp dir with a file in it
    temp_dir = tmp_path / "testdir"
    temp_dir.mkdir()
    (temp_dir / "afile.txt").write_text("hello")
    factory = DummyTmpPathFactory(
        basetemp=temp_dir, retention_policy="failed", given_basetemp=None
    )
    session = DummySession(factory)
    pytest_sessionfinish(session, exitstatus=0)  # 53.1μs -> 53.1μs (0.051% slower)


def test_policy_failed_exitstatus_nonzero_keeps_dir(tmp_path):
    """If exitstatus!=0, directory should not be removed."""
    temp_dir = tmp_path / "testdir"
    temp_dir.mkdir()
    (temp_dir / "afile.txt").write_text("hello")
    factory = DummyTmpPathFactory(
        basetemp=temp_dir, retention_policy="failed", given_basetemp=None
    )
    session = DummySession(factory)
    pytest_sessionfinish(session, exitstatus=1)  # 18.0μs -> 17.8μs (1.18% faster)


def test_policy_not_failed_keeps_dir(tmp_path):
    """If policy is not 'failed', directory should not be removed."""
    temp_dir = tmp_path / "testdir"
    temp_dir.mkdir()
    (temp_dir / "afile.txt").write_text("hello")
    factory = DummyTmpPathFactory(
        basetemp=temp_dir, retention_policy="all", given_basetemp=None
    )
    session = DummySession(factory)
    pytest_sessionfinish(session, exitstatus=0)  # 16.9μs -> 16.8μs (0.953% faster)


def test_given_basetemp_keeps_dir(tmp_path):
    """If _given_basetemp is not None, directory should not be removed."""
    temp_dir = tmp_path / "testdir"
    temp_dir.mkdir()
    (temp_dir / "afile.txt").write_text("hello")
    factory = DummyTmpPathFactory(
        basetemp=temp_dir,
        retention_policy="failed",
        given_basetemp=tmp_path / "something",
    )
    session = DummySession(factory)
    pytest_sessionfinish(session, exitstatus=0)  # 17.7μs -> 17.8μs (0.400% slower)


# EDGE TEST CASES


def test_basetemp_is_file(tmp_path):
    """If basetemp is a file, nothing should be removed and no error should be raised."""
    temp_file = tmp_path / "afile.txt"
    temp_file.write_text("hello")
    factory = DummyTmpPathFactory(
        basetemp=temp_file, retention_policy="failed", given_basetemp=None
    )
    session = DummySession(factory)
    pytest_sessionfinish(session, exitstatus=0)  # 6.46μs -> 6.48μs (0.293% slower)


def test_basetemp_dir_with_dead_symlink(tmp_path):
    """Dead symlinks in basetemp should be cleaned up, but not live symlinks."""
    temp_dir = tmp_path / "testdir"
    temp_dir.mkdir()
    # Create a dead symlink
    dead_target = tmp_path / "doesnotexist"
    dead_symlink = temp_dir / "deadlink"
    dead_symlink.symlink_to(dead_target)
    # Create a live symlink
    target_file = tmp_path / "target.txt"
    target_file.write_text("hi")
    live_symlink = temp_dir / "livelink"
    live_symlink.symlink_to(target_file)
    # Create a normal file
    (temp_dir / "afile.txt").write_text("hello")
    factory = DummyTmpPathFactory(
        basetemp=temp_dir, retention_policy="all", given_basetemp=None
    )
    session = DummySession(factory)
    pytest_sessionfinish(session, exitstatus=1)  # 118μs -> 43.0μs (175% faster)


def test_basetemp_is_empty_dir(tmp_path):
    """Empty basetemp should not error and should remain unless removal conditions met."""
    temp_dir = tmp_path / "testdir"
    temp_dir.mkdir()
    factory = DummyTmpPathFactory(
        basetemp=temp_dir, retention_policy="all", given_basetemp=None
    )
    session = DummySession(factory)
    pytest_sessionfinish(session, exitstatus=1)  # 13.1μs -> 12.8μs (2.33% faster)


def test_basetemp_dir_with_nested_dirs(tmp_path):
    """Nested directories should be removed if removal conditions met."""
    temp_dir = tmp_path / "testdir"
    nested_dir = temp_dir / "nested"
    nested_dir.mkdir(parents=True)
    (nested_dir / "afile.txt").write_text("hello")
    factory = DummyTmpPathFactory(
        basetemp=temp_dir, retention_policy="failed", given_basetemp=None
    )
    session = DummySession(factory)
    pytest_sessionfinish(session, exitstatus=0)  # 72.3μs -> 70.9μs (1.89% faster)


# LARGE SCALE TEST CASES


def test_large_number_of_files_and_symlinks(tmp_path):
    """Test performance and correctness with a large number of files and symlinks."""
    temp_dir = tmp_path / "testdir"
    temp_dir.mkdir()
    n_files = 200  # Not too large for test speed
    # Create files
    for i in range(n_files):
        (temp_dir / f"file_{i}.txt").write_text("data")
    # Create live symlinks
    for i in range(n_files):
        (temp_dir / f"live_symlink_{i}").symlink_to(temp_dir / f"file_{i}.txt")
    # Create dead symlinks
    for i in range(n_files):
        (temp_dir / f"dead_symlink_{i}").symlink_to(temp_dir / f"nonexistent_{i}.txt")
    factory = DummyTmpPathFactory(
        basetemp=temp_dir, retention_policy="all", given_basetemp=None
    )
    session = DummySession(factory)
    pytest_sessionfinish(session, exitstatus=1)  # 17.5ms -> 5.39ms (225% faster)
    # All dead symlinks should be gone, live ones and files should remain
    for i in range(n_files):
        pass


def test_large_scale_removal(tmp_path):
    """Test that a large directory is removed under the right conditions."""
    temp_dir = tmp_path / "testdir"
    temp_dir.mkdir()
    n_files = 500
    for i in range(n_files):
        (temp_dir / f"file_{i}.txt").write_text("data")
    factory = DummyTmpPathFactory(
        basetemp=temp_dir, retention_policy="failed", given_basetemp=None
    )
    session = DummySession(factory)
    pytest_sessionfinish(session, exitstatus=0)  # 3.44ms -> 3.42ms (0.416% faster)


def test_large_scale_keeps_on_wrong_policy(tmp_path):
    """Test that a large directory is NOT removed if policy is not 'failed'."""
    temp_dir = tmp_path / "testdir"
    temp_dir.mkdir()
    n_files = 500
    for i in range(n_files):
        (temp_dir / f"file_{i}.txt").write_text("data")
    factory = DummyTmpPathFactory(
        basetemp=temp_dir, retention_policy="all", given_basetemp=None
    )
    session = DummySession(factory)
    pytest_sessionfinish(session, exitstatus=0)  # 1.48ms -> 1.48ms (0.034% slower)


# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

from pathlib import Path
import shutil
import tempfile

# function to test
# imports
from _pytest.tmpdir import pytest_sessionfinish


# ---- UNIT TESTS ----


# Helper to create a fake session/config/tmp_path_factory structure
class FakeTmpPathFactory:
    def __init__(self, basetemp, retention_policy, given_basetemp):
        self._basetemp = basetemp
        self._retention_policy = retention_policy
        self._given_basetemp = given_basetemp


class FakeConfig:
    def __init__(self, tmp_path_factory):
        self._tmp_path_factory = tmp_path_factory


class FakeSession:
    def __init__(self, config):
        self.config = config


# 1. BASIC TEST CASES


def test_no_basetemp_returns(tmp_path_factory=None):
    """If basetemp is None, function should return immediately and do nothing."""
    tmp_path_factory = FakeTmpPathFactory(
        basetemp=None, retention_policy="failed", given_basetemp=None
    )
    config = FakeConfig(tmp_path_factory)
    session = FakeSession(config)
    # Should not raise or error
    pytest_sessionfinish(session, 0)  # 581ns -> 559ns (3.94% faster)


def test_policy_not_failed(tmp_path_factory=None):
    """If retention policy is not 'failed', directory should not be removed."""
    with tempfile.TemporaryDirectory() as tempdir:
        basetemp = Path(tempdir)
        # Create a file to check for removal
        (basetemp / "foo.txt").write_text("bar")
        tmp_path_factory = FakeTmpPathFactory(
            basetemp=basetemp, retention_policy="all", given_basetemp=None
        )
        config = FakeConfig(tmp_path_factory)
        session = FakeSession(config)
        pytest_sessionfinish(session, 0)  # 22.4μs -> 21.2μs (5.99% faster)


def test_exitstatus_not_zero(tmp_path_factory=None):
    """If exitstatus is not 0, directory should not be removed."""
    with tempfile.TemporaryDirectory() as tempdir:
        basetemp = Path(tempdir)
        (basetemp / "foo.txt").write_text("bar")
        tmp_path_factory = FakeTmpPathFactory(
            basetemp=basetemp, retention_policy="failed", given_basetemp=None
        )
        config = FakeConfig(tmp_path_factory)
        session = FakeSession(config)
        pytest_sessionfinish(session, 1)  # 20.1μs -> 19.5μs (3.34% faster)


def test_given_basetemp_is_set(tmp_path_factory=None):
    """If _given_basetemp is set, directory should not be removed."""
    with tempfile.TemporaryDirectory() as tempdir:
        basetemp = Path(tempdir)
        (basetemp / "foo.txt").write_text("bar")
        tmp_path_factory = FakeTmpPathFactory(
            basetemp=basetemp,
            retention_policy="failed",
            given_basetemp=Path("/tmp/something"),
        )
        config = FakeConfig(tmp_path_factory)
        session = FakeSession(config)
        pytest_sessionfinish(session, 0)  # 18.9μs -> 19.4μs (2.64% slower)


def test_removes_directory_on_successful_policy_failed(tmp_path_factory=None):
    """If exitstatus==0, policy=='failed', and _given_basetemp is None, directory should be removed."""
    with tempfile.TemporaryDirectory() as tempdir:
        basetemp = Path(tempdir)
        (basetemp / "foo.txt").write_text("bar")
        tmp_path_factory = FakeTmpPathFactory(
            basetemp=basetemp, retention_policy="failed", given_basetemp=None
        )
        config = FakeConfig(tmp_path_factory)
        session = FakeSession(config)
        pytest_sessionfinish(session, 0)  # 58.7μs -> 59.8μs (1.85% slower)


# 2. EDGE TEST CASES


def test_basetemp_is_file(tmp_path_factory=None):
    """If basetemp is a file, not a directory, should not attempt to remove or symlink clean."""
    with tempfile.NamedTemporaryFile() as temp_file:
        basetemp = Path(temp_file.name)
        tmp_path_factory = FakeTmpPathFactory(
            basetemp=basetemp, retention_policy="failed", given_basetemp=None
        )
        config = FakeConfig(tmp_path_factory)
        session = FakeSession(config)
        # Should not raise
        pytest_sessionfinish(session, 0)  # 7.93μs -> 8.64μs (8.24% slower)


def test_basetemp_is_empty_dir(tmp_path_factory=None):
    """If basetemp is an empty directory, should remove it if conditions are met."""
    with tempfile.TemporaryDirectory() as tempdir:
        basetemp = Path(tempdir)
        tmp_path_factory = FakeTmpPathFactory(
            basetemp=basetemp, retention_policy="failed", given_basetemp=None
        )
        config = FakeConfig(tmp_path_factory)
        session = FakeSession(config)
        pytest_sessionfinish(session, 0)  # 48.9μs -> 49.8μs (1.74% slower)


def test_removes_dead_symlinks():
    """Should remove dead symlinks in basetemp directory, but not valid ones."""
    with tempfile.TemporaryDirectory() as tempdir:
        basetemp = Path(tempdir)
        # Valid symlink
        real_file = basetemp / "real.txt"
        real_file.write_text("hi")
        valid_symlink = basetemp / "valid_link"
        valid_symlink.symlink_to(real_file)
        # Dead symlink
        dead_symlink = basetemp / "dead_link"
        dead_symlink.symlink_to(basetemp / "doesnotexist.txt")
        # Should not trigger rmtree (simulate: exitstatus!=0)
        tmp_path_factory = FakeTmpPathFactory(
            basetemp=basetemp, retention_policy="failed", given_basetemp=None
        )
        config = FakeConfig(tmp_path_factory)
        session = FakeSession(config)
        pytest_sessionfinish(session, 1)  # 99.2μs -> 43.7μs (127% faster)


def test_symlink_to_dir_is_dead():
    """Should remove dead symlink to a directory."""
    with tempfile.TemporaryDirectory() as tempdir:
        basetemp = Path(tempdir)
        # Create a subdir and a symlink to it, then remove subdir
        subdir = basetemp / "subdir"
        subdir.mkdir()
        dir_symlink = basetemp / "dir_link"
        dir_symlink.symlink_to(subdir)
        # Remove the target directory to make the symlink dead
        shutil.rmtree(subdir)
        # Should not trigger rmtree (simulate: exitstatus!=0)
        tmp_path_factory = FakeTmpPathFactory(
            basetemp=basetemp, retention_policy="failed", given_basetemp=None
        )
        config = FakeConfig(tmp_path_factory)
        session = FakeSession(config)
        pytest_sessionfinish(session, 1)  # 57.6μs -> 26.2μs (120% faster)


def test_symlink_to_file_is_valid():
    """Should not remove valid symlink to a file."""
    with tempfile.TemporaryDirectory() as tempdir:
        basetemp = Path(tempdir)
        real_file = basetemp / "real.txt"
        real_file.write_text("hi")
        valid_symlink = basetemp / "valid_link"
        valid_symlink.symlink_to(real_file)
        tmp_path_factory = FakeTmpPathFactory(
            basetemp=basetemp, retention_policy="failed", given_basetemp=None
        )
        config = FakeConfig(tmp_path_factory)
        session = FakeSession(config)
        pytest_sessionfinish(session, 1)  # 53.5μs -> 24.1μs (122% faster)


def test_directory_with_permissions_issue(monkeypatch):
    """If rmtree fails due to permissions, should ignore error and not raise."""
    with tempfile.TemporaryDirectory() as tempdir:
        basetemp = Path(tempdir)
        (basetemp / "foo.txt").write_text("bar")
        tmp_path_factory = FakeTmpPathFactory(
            basetemp=basetemp, retention_policy="failed", given_basetemp=None
        )
        config = FakeConfig(tmp_path_factory)
        session = FakeSession(config)
        # Monkeypatch rmtree to raise exception
        called = {}

        def fake_rmtree(path, ignore_errors):
            called["yes"] = True
            raise PermissionError("Fake permission error")

        monkeypatch.setattr(shutil, "rmtree", fake_rmtree)
        # Should not raise
        pytest_sessionfinish(session, 0)  # 58.2μs -> 59.8μs (2.74% slower)


# 3. LARGE SCALE TEST CASES


def test_large_number_of_files_and_symlinks():
    """Test cleanup with a large number of files and symlinks."""
    with tempfile.TemporaryDirectory() as tempdir:
        basetemp = Path(tempdir)
        # Create 500 files and 500 symlinks (half valid, half dead)
        files = []
        for i in range(250):
            f = basetemp / f"file_{i}.txt"
            f.write_text("data")
            files.append(f)
        # Valid symlinks
        for i in range(250):
            (basetemp / f"valid_link_{i}").symlink_to(files[i])
        # Dead symlinks
        for i in range(250, 500):
            (basetemp / f"dead_link_{i}").symlink_to(basetemp / f"missing_{i}.txt")
        # Should not remove directory (exitstatus!=0)
        tmp_path_factory = FakeTmpPathFactory(
            basetemp=basetemp, retention_policy="failed", given_basetemp=None
        )
        config = FakeConfig(tmp_path_factory)
        session = FakeSession(config)
        pytest_sessionfinish(session, 1)  # 14.7ms -> 5.62ms (162% faster)
        # All dead symlinks should be gone
        for i in range(250, 500):
            pass
        # All valid symlinks should remain
        for i in range(250):
            link = basetemp / f"valid_link_{i}"


def test_large_directory_is_removed_on_success():
    """Test that a large directory is removed if all conditions are met."""
    with tempfile.TemporaryDirectory() as tempdir:
        basetemp = Path(tempdir)
        # Create 100 files
        for i in range(100):
            (basetemp / f"file_{i}.txt").write_text("data")
        tmp_path_factory = FakeTmpPathFactory(
            basetemp=basetemp, retention_policy="failed", given_basetemp=None
        )
        config = FakeConfig(tmp_path_factory)
        session = FakeSession(config)
        pytest_sessionfinish(session, 0)  # 698μs -> 697μs (0.038% faster)


# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

To edit these changes git checkout codeflash/optimize-pytest_sessionfinish-mi9jubn0 and push.