⚡️ Speed up function is_pr_draft by 35% in PR #322 (fix/skip-optimization-for-draft-prs)

[codeflash-ai]

⚡️ This pull request contains optimizations for PR #322

If you approve this dependent PR, these changes will be merged into the original PR branch fix/skip-optimization-for-draft-prs.

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

---

📄 35% (0.35x) speedup for is_pr_draft in codeflash/optimization/optimizer.py

⏱️ Runtime : 6.43 milliseconds → 4.78 milliseconds (best of 92 runs)

📝 Explanation and details

Here’s a faster, more memory-efficient rewrite of your program. Optimizations include.

• Removed duplicate get_pr_number import and moved/corrected its singleton responsibility.
• Used direct, local variable access rather than repeatedly referencing imported modules.
• Used os.environ[] instead of os.environ.get for critical env lookups in a try-block (avoids lookup cost when you know failure will land in except anyway).
• Used direct file open for reading (avoiding Path overhead).
• Avoided reading/parsing JSON and dictionary keys if the PR number/env is missing.
• Reduced exception handling scope to only JSON/file-related operations.

All existing comments are preserved except where the code was made more efficient.

Key points:

• File reading is done only if both env vars are present, before JSON parsing.
• The exception only wraps file I/O + JSON parsing, not the env checks, so it's tighter/faster in normal runs.
• No Path. Used built-in open for speed.
• Early returns for failures, so the function does the minimum work necessary.
• Single access to environment variables (no redundancy).

Functional output is preserved.

✅ Correctness verification report:

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

🌀 Generated Regression Tests Details

from __future__ import annotations

import json
import os
import tempfile
from functools import lru_cache
from pathlib import Path
from typing import Optional

# imports
import pytest  # used for our unit tests
from codeflash.cli_cmds.console import logger
from codeflash.code_utils.env_utils import get_pr_number
from codeflash.optimization.optimizer import is_pr_draft

# --- Basic Test Cases ---

def test_draft_true(monkeypatch, tmp_path):
    """Test: PR is draft, all environment variables present."""
    # Prepare event JSON file
    event_data = {"pull_request": {"draft": True}}
    event_file = tmp_path / "event.json"
    event_file.write_text(json.dumps(event_data))
    # Set env variables
    monkeypatch.setenv("CODEFLASH_PR_NUMBER", "42")
    monkeypatch.setenv("GITHUB_EVENT_PATH", str(event_file))
    # Should return True
    codeflash_output = is_pr_draft() # 47.8μs -> 35.2μs

def test_draft_false(monkeypatch, tmp_path):
    """Test: PR is not draft, all environment variables present."""
    event_data = {"pull_request": {"draft": False}}
    event_file = tmp_path / "event.json"
    event_file.write_text(json.dumps(event_data))
    monkeypatch.setenv("CODEFLASH_PR_NUMBER", "7")
    monkeypatch.setenv("GITHUB_EVENT_PATH", str(event_file))
    codeflash_output = is_pr_draft() # 47.2μs -> 35.0μs

def test_no_pr_number(monkeypatch, tmp_path):
    """Test: CODEFLASH_PR_NUMBER missing, should return False."""
    event_data = {"pull_request": {"draft": True}}
    event_file = tmp_path / "event.json"
    event_file.write_text(json.dumps(event_data))
    monkeypatch.delenv("CODEFLASH_PR_NUMBER", raising=False)
    monkeypatch.setenv("GITHUB_EVENT_PATH", str(event_file))
    codeflash_output = is_pr_draft() # 5.47μs -> 5.13μs

def test_no_event_path(monkeypatch):
    """Test: GITHUB_EVENT_PATH missing, should return False."""
    monkeypatch.setenv("CODEFLASH_PR_NUMBER", "123")
    monkeypatch.delenv("GITHUB_EVENT_PATH", raising=False)
    codeflash_output = is_pr_draft() # 4.95μs -> 4.53μs

def test_both_env_missing(monkeypatch):
    """Test: Both env vars missing, should return False."""
    monkeypatch.delenv("CODEFLASH_PR_NUMBER", raising=False)
    monkeypatch.delenv("GITHUB_EVENT_PATH", raising=False)
    codeflash_output = is_pr_draft() # 4.79μs -> 4.68μs

# --- Edge Test Cases ---

def test_event_path_file_missing(monkeypatch, tmp_path):
    """Test: GITHUB_EVENT_PATH points to non-existent file, should return False."""
    non_existent = tmp_path / "does_not_exist.json"
    monkeypatch.setenv("CODEFLASH_PR_NUMBER", "5")
    monkeypatch.setenv("GITHUB_EVENT_PATH", str(non_existent))
    codeflash_output = is_pr_draft() # 101μs -> 98.3μs

def test_event_path_not_json(monkeypatch, tmp_path):
    """Test: GITHUB_EVENT_PATH points to invalid JSON file, should return False."""
    event_file = tmp_path / "event.json"
    event_file.write_text("not a json")
    monkeypatch.setenv("CODEFLASH_PR_NUMBER", "5")
    monkeypatch.setenv("GITHUB_EVENT_PATH", str(event_file))
    codeflash_output = is_pr_draft() # 117μs -> 104μs

def test_event_path_no_pull_request_key(monkeypatch, tmp_path):
    """Test: JSON missing 'pull_request' key, should return False."""
    event_data = {"not_pull_request": {"draft": True}}
    event_file = tmp_path / "event.json"
    event_file.write_text(json.dumps(event_data))
    monkeypatch.setenv("CODEFLASH_PR_NUMBER", "5")
    monkeypatch.setenv("GITHUB_EVENT_PATH", str(event_file))
    codeflash_output = is_pr_draft() # 110μs -> 94.9μs

def test_event_path_pull_request_no_draft_key(monkeypatch, tmp_path):
    """Test: JSON missing 'draft' key in 'pull_request', should return False."""
    event_data = {"pull_request": {"not_draft": True}}
    event_file = tmp_path / "event.json"
    event_file.write_text(json.dumps(event_data))
    monkeypatch.setenv("CODEFLASH_PR_NUMBER", "5")
    monkeypatch.setenv("GITHUB_EVENT_PATH", str(event_file))
    codeflash_output = is_pr_draft() # 111μs -> 92.2μs

def test_codeflash_pr_number_non_integer(monkeypatch, tmp_path):
    """Test: CODEFLASH_PR_NUMBER is not an integer, should return False."""
    event_data = {"pull_request": {"draft": True}}
    event_file = tmp_path / "event.json"
    event_file.write_text(json.dumps(event_data))
    monkeypatch.setenv("CODEFLASH_PR_NUMBER", "not_an_int")
    monkeypatch.setenv("GITHUB_EVENT_PATH", str(event_file))
    codeflash_output = is_pr_draft() # 66.3μs -> 7.31μs

def test_codeflash_pr_number_zero(monkeypatch, tmp_path):
    """Test: CODEFLASH_PR_NUMBER is zero (valid int), should still work."""
    event_data = {"pull_request": {"draft": True}}
    event_file = tmp_path / "event.json"
    event_file.write_text(json.dumps(event_data))
    monkeypatch.setenv("CODEFLASH_PR_NUMBER", "0")
    monkeypatch.setenv("GITHUB_EVENT_PATH", str(event_file))
    codeflash_output = is_pr_draft() # 48.2μs -> 35.9μs

def test_codeflash_pr_number_negative(monkeypatch, tmp_path):
    """Test: CODEFLASH_PR_NUMBER is negative (valid int), should still work."""
    event_data = {"pull_request": {"draft": False}}
    event_file = tmp_path / "event.json"
    event_file.write_text(json.dumps(event_data))
    monkeypatch.setenv("CODEFLASH_PR_NUMBER", "-1")
    monkeypatch.setenv("GITHUB_EVENT_PATH", str(event_file))
    codeflash_output = is_pr_draft() # 47.5μs -> 35.2μs

def test_event_path_is_directory(monkeypatch, tmp_path):
    """Test: GITHUB_EVENT_PATH points to a directory, not a file, should return False."""
    monkeypatch.setenv("CODEFLASH_PR_NUMBER", "123")
    monkeypatch.setenv("GITHUB_EVENT_PATH", str(tmp_path))
    codeflash_output = is_pr_draft() # 94.2μs -> 76.6μs

def test_event_path_empty_file(monkeypatch, tmp_path):
    """Test: GITHUB_EVENT_PATH points to an empty file, should return False."""
    event_file = tmp_path / "empty.json"
    event_file.write_text("")
    monkeypatch.setenv("CODEFLASH_PR_NUMBER", "123")
    monkeypatch.setenv("GITHUB_EVENT_PATH", str(event_file))
    codeflash_output = is_pr_draft() # 115μs -> 98.9μs

def test_event_path_extra_keys(monkeypatch, tmp_path):
    """Test: JSON has extra keys, but correct structure, should work."""
    event_data = {
        "pull_request": {"draft": True, "other": 1},
        "extra": 42
    }
    event_file = tmp_path / "event.json"
    event_file.write_text(json.dumps(event_data))
    monkeypatch.setenv("CODEFLASH_PR_NUMBER", "77")
    monkeypatch.setenv("GITHUB_EVENT_PATH", str(event_file))
    codeflash_output = is_pr_draft() # 49.1μs -> 36.4μs

def test_draft_value_is_none(monkeypatch, tmp_path):
    """Test: 'draft' key is present but value is None, should return False."""
    event_data = {"pull_request": {"draft": None}}
    event_file = tmp_path / "event.json"
    event_file.write_text(json.dumps(event_data))
    monkeypatch.setenv("CODEFLASH_PR_NUMBER", "88")
    monkeypatch.setenv("GITHUB_EVENT_PATH", str(event_file))
    codeflash_output = is_pr_draft() # 47.5μs -> 35.2μs

def test_draft_value_is_string_true(monkeypatch, tmp_path):
    """Test: 'draft' key is string 'true', not boolean, should raise KeyError and return False."""
    event_data = {"pull_request": {"draft": "true"}}
    event_file = tmp_path / "event.json"
    event_file.write_text(json.dumps(event_data))
    monkeypatch.setenv("CODEFLASH_PR_NUMBER", "99")
    monkeypatch.setenv("GITHUB_EVENT_PATH", str(event_file))
    # Should return "true" (truthy) but function expects boolean, so will return "true"
    # However, since the code does not type check, it will return "true" (a string), which is truthy.
    # But the function's contract is to return bool. Let's check for bool output.
    codeflash_output = is_pr_draft(); result = codeflash_output # 47.5μs -> 34.7μs

def test_draft_value_is_integer(monkeypatch, tmp_path):
    """Test: 'draft' key is integer 1, not boolean, should return 1 (truthy)."""
    event_data = {"pull_request": {"draft": 1}}
    event_file = tmp_path / "event.json"
    event_file.write_text(json.dumps(event_data))
    monkeypatch.setenv("CODEFLASH_PR_NUMBER", "100")
    monkeypatch.setenv("GITHUB_EVENT_PATH", str(event_file))
    codeflash_output = is_pr_draft(); result = codeflash_output # 47.1μs -> 35.2μs

# --- Large Scale Test Cases ---

def test_large_json_file(monkeypatch, tmp_path):
    """Test: Large JSON file with many keys, correct draft value."""
    # Create a large dict with 999 dummy keys and the real one
    large_data = {f"key_{i}": i for i in range(999)}
    large_data["pull_request"] = {"draft": True}
    event_file = tmp_path / "event.json"
    event_file.write_text(json.dumps(large_data))
    monkeypatch.setenv("CODEFLASH_PR_NUMBER", "1000")
    monkeypatch.setenv("GITHUB_EVENT_PATH", str(event_file))
    codeflash_output = is_pr_draft() # 197μs -> 184μs

def test_many_calls_with_different_env(monkeypatch, tmp_path):
    """Test: Multiple calls with different envs, ensure lru_cache does not leak state."""
    # First, set up a draft PR
    event_data1 = {"pull_request": {"draft": True}}
    event_file1 = tmp_path / "event1.json"
    event_file1.write_text(json.dumps(event_data1))
    # Second, set up a non-draft PR
    event_data2 = {"pull_request": {"draft": False}}
    event_file2 = tmp_path / "event2.json"
    event_file2.write_text(json.dumps(event_data2))

    # Call with first PR
    monkeypatch.setenv("CODEFLASH_PR_NUMBER", "111")
    monkeypatch.setenv("GITHUB_EVENT_PATH", str(event_file1))
    get_pr_number.cache_clear()
    codeflash_output = is_pr_draft() # 37.1μs -> 26.1μs

    # Call with second PR
    monkeypatch.setenv("CODEFLASH_PR_NUMBER", "222")
    monkeypatch.setenv("GITHUB_EVENT_PATH", str(event_file2))
    get_pr_number.cache_clear()
    codeflash_output = is_pr_draft() # 37.1μs -> 26.1μs

def test_performance_large_number_of_calls(monkeypatch, tmp_path):
    """Test: Call is_pr_draft 100 times with valid PR, ensure always True and no exceptions."""
    event_data = {"pull_request": {"draft": True}}
    event_file = tmp_path / "event.json"
    event_file.write_text(json.dumps(event_data))
    monkeypatch.setenv("CODEFLASH_PR_NUMBER", "333")
    monkeypatch.setenv("GITHUB_EVENT_PATH", str(event_file))
    for _ in range(100):
        get_pr_number.cache_clear()
        codeflash_output = is_pr_draft() # 24.6μs -> 16.2μs

def test_large_json_with_deeply_nested(monkeypatch, tmp_path):
    """Test: Large JSON with deeply nested irrelevant data, but correct draft value."""
    # Add a deep structure
    deep = {}
    cur = deep
    for i in range(500):
        cur[f"level_{i}"] = {}
        cur = cur[f"level_{i}"]
    # Add the real key
    event_data = {"pull_request": {"draft": False}, "deep": deep}
    event_file = tmp_path / "event.json"
    event_file.write_text(json.dumps(event_data))
    monkeypatch.setenv("CODEFLASH_PR_NUMBER", "444")
    monkeypatch.setenv("GITHUB_EVENT_PATH", str(event_file))
    codeflash_output = is_pr_draft() # 130μs -> 112μs
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

from __future__ import annotations

import json
import os
import tempfile
from functools import lru_cache
from pathlib import Path
from typing import Optional

# imports
import pytest  # used for our unit tests
from codeflash.cli_cmds.console import logger
from codeflash.code_utils.env_utils import get_pr_number
from codeflash.optimization.optimizer import is_pr_draft


def write_event_file(tmp_path, event_data):
    # Helper to write a JSON event file
    event_file = tmp_path / "event.json"
    with open(event_file, "w") as f:
        json.dump(event_data, f)
    return str(event_file)

# 1. Basic Test Cases

def test_draft_true(monkeypatch, tmp_path):
    # PR is draft: True
    monkeypatch.setenv("CODEFLASH_PR_NUMBER", "123")
    event_data = {"pull_request": {"draft": True}}
    event_file = write_event_file(tmp_path, event_data)
    monkeypatch.setenv("GITHUB_EVENT_PATH", event_file)
    codeflash_output = is_pr_draft() # 47.8μs -> 35.2μs

def test_draft_false(monkeypatch, tmp_path):
    # PR is draft: False
    monkeypatch.setenv("CODEFLASH_PR_NUMBER", "456")
    event_data = {"pull_request": {"draft": False}}
    event_file = write_event_file(tmp_path, event_data)
    monkeypatch.setenv("GITHUB_EVENT_PATH", event_file)
    codeflash_output = is_pr_draft() # 47.2μs -> 35.0μs

def test_no_pr_number(monkeypatch, tmp_path):
    # No CODEFLASH_PR_NUMBER set, should return False
    event_data = {"pull_request": {"draft": True}}
    event_file = write_event_file(tmp_path, event_data)
    monkeypatch.setenv("GITHUB_EVENT_PATH", event_file)
    codeflash_output = is_pr_draft() # 5.47μs -> 5.13μs

def test_no_event_path(monkeypatch):
    # No GITHUB_EVENT_PATH set, should return False
    monkeypatch.setenv("CODEFLASH_PR_NUMBER", "123")
    codeflash_output = is_pr_draft() # 4.95μs -> 4.53μs

# 2. Edge Test Cases

def test_event_file_missing(monkeypatch, tmp_path):
    # GITHUB_EVENT_PATH points to a non-existent file
    monkeypatch.setenv("CODEFLASH_PR_NUMBER", "789")
    missing_file = str(tmp_path / "does_not_exist.json")
    monkeypatch.setenv("GITHUB_EVENT_PATH", missing_file)
    codeflash_output = is_pr_draft() # 112μs -> 101μs

def test_event_file_invalid_json(monkeypatch, tmp_path):
    # GITHUB_EVENT_PATH points to a file with invalid JSON
    monkeypatch.setenv("CODEFLASH_PR_NUMBER", "789")
    event_file = tmp_path / "bad.json"
    with open(event_file, "w") as f:
        f.write("{not valid json}")
    monkeypatch.setenv("GITHUB_EVENT_PATH", str(event_file))
    codeflash_output = is_pr_draft() # 117μs -> 104μs

def test_event_file_missing_pull_request_key(monkeypatch, tmp_path):
    # JSON file missing 'pull_request' key
    monkeypatch.setenv("CODEFLASH_PR_NUMBER", "789")
    event_data = {"not_pull_request": {"draft": True}}
    event_file = write_event_file(tmp_path, event_data)
    monkeypatch.setenv("GITHUB_EVENT_PATH", event_file)
    codeflash_output = is_pr_draft() # 108μs -> 92.7μs

def test_event_file_missing_draft_key(monkeypatch, tmp_path):
    # JSON file missing 'draft' key inside 'pull_request'
    monkeypatch.setenv("CODEFLASH_PR_NUMBER", "789")
    event_data = {"pull_request": {"not_draft": True}}
    event_file = write_event_file(tmp_path, event_data)
    monkeypatch.setenv("GITHUB_EVENT_PATH", event_file)
    codeflash_output = is_pr_draft() # 107μs -> 92.4μs

def test_pr_number_not_int(monkeypatch, tmp_path):
    # CODEFLASH_PR_NUMBER is not an integer
    monkeypatch.setenv("CODEFLASH_PR_NUMBER", "not_an_int")
    event_data = {"pull_request": {"draft": True}}
    event_file = write_event_file(tmp_path, event_data)
    monkeypatch.setenv("GITHUB_EVENT_PATH", event_file)
    codeflash_output = is_pr_draft() # 67.8μs -> 7.14μs

def test_event_file_draft_null(monkeypatch, tmp_path):
    # 'draft' key is None
    monkeypatch.setenv("CODEFLASH_PR_NUMBER", "123")
    event_data = {"pull_request": {"draft": None}}
    event_file = write_event_file(tmp_path, event_data)
    monkeypatch.setenv("GITHUB_EVENT_PATH", event_file)
    codeflash_output = is_pr_draft() # 47.5μs -> 35.8μs

def test_event_file_draft_string(monkeypatch, tmp_path):
    # 'draft' key is a string "true"
    monkeypatch.setenv("CODEFLASH_PR_NUMBER", "123")
    event_data = {"pull_request": {"draft": "true"}}
    event_file = write_event_file(tmp_path, event_data)
    monkeypatch.setenv("GITHUB_EVENT_PATH", event_file)
    # Should return "true", not True
    codeflash_output = is_pr_draft() # 47.2μs -> 35.0μs

def test_event_file_extra_keys(monkeypatch, tmp_path):
    # Extra keys in JSON, should not affect result
    monkeypatch.setenv("CODEFLASH_PR_NUMBER", "123")
    event_data = {
        "pull_request": {"draft": True, "other": 1},
        "extra": {"foo": "bar"}
    }
    event_file = write_event_file(tmp_path, event_data)
    monkeypatch.setenv("GITHUB_EVENT_PATH", event_file)
    codeflash_output = is_pr_draft() # 48.5μs -> 36.1μs

def test_event_file_minimal(monkeypatch, tmp_path):
    # Minimal valid JSON
    monkeypatch.setenv("CODEFLASH_PR_NUMBER", "1")
    event_data = {"pull_request": {"draft": False}}
    event_file = write_event_file(tmp_path, event_data)
    monkeypatch.setenv("GITHUB_EVENT_PATH", event_file)
    codeflash_output = is_pr_draft() # 46.9μs -> 35.2μs

def test_event_file_empty(monkeypatch, tmp_path):
    # Empty JSON file
    monkeypatch.setenv("CODEFLASH_PR_NUMBER", "1")
    event_file = tmp_path / "empty.json"
    with open(event_file, "w") as f:
        f.write("{}")
    monkeypatch.setenv("GITHUB_EVENT_PATH", str(event_file))
    codeflash_output = is_pr_draft() # 111μs -> 97.4μs

def test_event_file_not_json(monkeypatch, tmp_path):
    # File is not JSON (e.g., plain text)
    monkeypatch.setenv("CODEFLASH_PR_NUMBER", "1")
    event_file = tmp_path / "notjson.txt"
    with open(event_file, "w") as f:
        f.write("hello world")
    monkeypatch.setenv("GITHUB_EVENT_PATH", str(event_file))
    codeflash_output = is_pr_draft() # 118μs -> 101μs

def test_event_file_permissions(monkeypatch, tmp_path):
    # File exists but is not readable (simulate by removing permissions)
    monkeypatch.setenv("CODEFLASH_PR_NUMBER", "1")
    event_data = {"pull_request": {"draft": True}}
    event_file = write_event_file(tmp_path, event_data)
    os.chmod(event_file, 0)  # Remove all permissions
    monkeypatch.setenv("GITHUB_EVENT_PATH", event_file)
    try:
        codeflash_output = is_pr_draft() # 86.7μs -> 69.4μs
    finally:
        os.chmod(event_file, 0o600)  # Restore permissions for cleanup

# 3. Large Scale Test Cases

def test_large_json_file(monkeypatch, tmp_path):
    # Large JSON file with many keys, but correct pull_request/draft
    monkeypatch.setenv("CODEFLASH_PR_NUMBER", "999")
    # Create a large dict with irrelevant data
    large_data = {f"key_{i}": i for i in range(900)}
    large_data["pull_request"] = {"draft": True}
    event_file = write_event_file(tmp_path, large_data)
    monkeypatch.setenv("GITHUB_EVENT_PATH", event_file)
    codeflash_output = is_pr_draft() # 197μs -> 184μs

def test_multiple_calls_cache(monkeypatch, tmp_path):
    # Test that lru_cache does not break correct behavior across calls
    monkeypatch.setenv("CODEFLASH_PR_NUMBER", "100")
    event_data = {"pull_request": {"draft": True}}
    event_file = write_event_file(tmp_path, event_data)
    monkeypatch.setenv("GITHUB_EVENT_PATH", event_file)
    # First call
    codeflash_output = is_pr_draft() # 44.3μs -> 31.5μs
    # Change event file to draft: False, should still return True due to lru_cache on get_pr_number
    event_data2 = {"pull_request": {"draft": False}}
    event_file2 = write_event_file(tmp_path, event_data2)
    monkeypatch.setenv("GITHUB_EVENT_PATH", event_file2)
    codeflash_output = is_pr_draft() # 44.3μs -> 31.5μs

def test_many_invocations(monkeypatch, tmp_path):
    # Simulate many invocations with different PR numbers and event files
    for i in range(10):
        monkeypatch.setenv("CODEFLASH_PR_NUMBER", str(i))
        event_data = {"pull_request": {"draft": i % 2 == 0}}
        event_file = write_event_file(tmp_path, event_data)
        monkeypatch.setenv("GITHUB_EVENT_PATH", event_file)
        get_pr_number.cache_clear()
        expected = (i % 2 == 0)
        codeflash_output = is_pr_draft() # 41.8μs -> 31.0μs

def test_large_number_of_env_vars(monkeypatch, tmp_path):
    # Set many unrelated env vars, ensure function still works
    for i in range(500):
        monkeypatch.setenv(f"UNRELATED_ENV_{i}", str(i))
    monkeypatch.setenv("CODEFLASH_PR_NUMBER", "888")
    event_data = {"pull_request": {"draft": False}}
    event_file = write_event_file(tmp_path, event_data)
    monkeypatch.setenv("GITHUB_EVENT_PATH", event_file)
    codeflash_output = is_pr_draft() # 48.2μs -> 35.8μs
# 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-pr322-2025-06-16T16.36.11 and push.