⚡️ Speed up method EnvironmentReader.use by 193%

[codeflash-ai]

📄 193% (1.93x) speedup for EnvironmentReader.use in graphrag/config/environment_reader.py

⏱️ Runtime : 732 microseconds → 250 microseconds (best of 292 runs)

📝 Explanation and details

The optimization achieves a 193% speedup by eliminating expensive closure creation and function definition overhead on every call to use().

Key Change: The context manager function was moved from being defined inside the use() method to module scope as _config_context(), taking the stack and value as parameters.

Why This Is Faster:

• Original code: Every call to use() creates a new closure (config_context) with access to self, requiring Python to allocate memory for the closure object and set up variable bindings. The line profiler shows 73.5% of time spent in function definition.
• Optimized code: Uses a pre-defined module-level function, eliminating closure creation overhead entirely. The stack is passed as a parameter instead of being captured.

Performance Impact by Test Type:

• Basic operations (single context usage): ~3x faster due to eliminated closure overhead
• Nested contexts: Even better gains since the optimization compounds with multiple calls
• Large-scale tests (1000+ iterations): Maximum benefit as closure creation costs accumulate

This optimization is particularly effective for high-frequency usage patterns where use() is called repeatedly, as each call no longer pays the cost of creating and destroying closure objects.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

from contextlib import contextmanager
from typing import Any

# imports
import pytest  # used for our unit tests
from environs import Env
from graphrag.config.environment_reader import EnvironmentReader

# unit tests

# Helper class to mock Env, since we don't use its features
class DummyEnv(Env):
    pass

@pytest.fixture
def env_reader():
    """Fixture to create a fresh EnvironmentReader for each test."""
    return EnvironmentReader(DummyEnv())

# 1. Basic Test Cases

def test_use_pushes_and_pops_dict(env_reader):
    """Test that use pushes a dict and pops it after context."""
    value = {'a': 1, 'b': 2}
    with env_reader.use(value):
        pass

def test_use_with_none_pushes_empty_dict(env_reader):
    """Test that use(None) pushes an empty dict."""
    with env_reader.use(None):
        pass

def test_use_with_empty_dict(env_reader):
    """Test that use({}) pushes an empty dict."""
    with env_reader.use({}):
        pass

def test_use_with_non_dict_value(env_reader):
    """Test that use with a non-dict value pushes that value."""
    value = [1, 2, 3]
    with env_reader.use(value):
        pass

def test_use_with_integer(env_reader):
    """Test that use with an integer pushes the integer."""
    value = 42
    with env_reader.use(value):
        pass

def test_use_with_string(env_reader):
    """Test that use with a string pushes the string."""
    value = "hello"
    with env_reader.use(value):
        pass

# 2. Edge Test Cases

def test_use_nested_contexts(env_reader):
    """Test nested use contexts and stack behavior."""
    val1 = {'x': 1}
    val2 = {'y': 2}
    with env_reader.use(val1):
        with env_reader.use(val2):
            pass

def test_use_with_falsey_values(env_reader):
    """Test use with various falsey values."""
    for falsey in [None, False, 0, '', [], {}, set()]:
        with env_reader.use(falsey):
            # None should push {}, others should push themselves
            expected = {} if falsey is None else falsey

def test_use_context_exception_handling(env_reader):
    """Test that config_stack is popped even if an exception occurs in context."""
    value = {'err': True}
    try:
        with env_reader.use(value):
            raise RuntimeError("Test Exception")
    except RuntimeError:
        pass

def test_use_multiple_sequential_contexts(env_reader):
    """Test multiple sequential contexts."""
    values = [{'a': 1}, {'b': 2}, {'c': 3}]
    for val in values:
        with env_reader.use(val):
            pass

def test_use_stack_integrity_with_manual_pop(env_reader):
    """Test stack integrity if config_stack is manually popped inside context."""
    value = {'manual': True}
    with env_reader.use(value):
        env_reader._config_stack.pop()
        # After context exit, pop will fail (should raise IndexError)
        with pytest.raises(IndexError):
            # This will be triggered by finally in use
            pass  # Context exit triggers the pop

def test_use_with_object_reference(env_reader):
    """Test that use pushes object references correctly."""
    class Dummy:
        pass
    obj = Dummy()
    with env_reader.use(obj):
        pass

# 3. Large Scale Test Cases


def test_use_with_large_dict(env_reader):
    """Test pushing a large dict into the stack."""
    large_dict = {str(i): i for i in range(1000)}
    with env_reader.use(large_dict):
        pass

def test_use_with_large_list(env_reader):
    """Test pushing a large list into the stack."""
    large_list = list(range(1000))
    with env_reader.use(large_list):
        pass

def test_use_with_large_string(env_reader):
    """Test pushing a large string into the stack."""
    large_string = "x" * 1000
    with env_reader.use(large_string):
        pass

def test_use_with_large_set(env_reader):
    """Test pushing a large set into the stack."""
    large_set = set(range(1000))
    with env_reader.use(large_set):
        pass

def test_use_with_large_sequential_contexts(env_reader):
    """Test many sequential contexts with large values."""
    for i in range(1000):
        val = {'index': i}
        with env_reader.use(val):
            pass
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
from contextlib import contextmanager
from typing import Any

# imports
import pytest  # used for our unit tests
from environs import Env
from graphrag.config.environment_reader import EnvironmentReader

# unit tests

# Helper: Dummy Env class for testing (since we don't use Env in our tests)
class DummyEnv:
    pass

# --------- BASIC TEST CASES ---------

def test_use_pushes_dict_to_stack():
    """Test that use pushes a dict to the config stack and pops it after context."""
    reader = EnvironmentReader(DummyEnv())
    test_dict = {"foo": "bar"}
    with reader.use(test_dict):
        pass

def test_use_pushes_none_as_empty_dict():
    """Test that use(None) pushes an empty dict to the config stack."""
    reader = EnvironmentReader(DummyEnv())
    with reader.use(None):
        pass

def test_use_pushes_empty_dict():
    """Test that use({}) pushes an empty dict to the config stack."""
    reader = EnvironmentReader(DummyEnv())
    with reader.use({}):
        pass

def test_use_pushes_list_object():
    """Test that use with a non-dict object (e.g., list) pushes it as-is."""
    reader = EnvironmentReader(DummyEnv())
    test_list = [1, 2, 3]
    with reader.use(test_list):
        pass

def test_use_pushes_integer_object():
    """Test that use with an int pushes it as-is."""
    reader = EnvironmentReader(DummyEnv())
    test_int = 42
    with reader.use(test_int):
        pass

def test_use_pushes_string_object():
    """Test that use with a string pushes it as-is."""
    reader = EnvironmentReader(DummyEnv())
    test_str = "hello"
    with reader.use(test_str):
        pass

# --------- EDGE TEST CASES ---------

def test_use_nested_contexts():
    """Test that nested use contexts stack values correctly."""
    reader = EnvironmentReader(DummyEnv())
    d1 = {"a": 1}
    d2 = {"b": 2}
    with reader.use(d1):
        with reader.use(d2):
            pass

def test_use_nested_none_and_dict():
    """Test nested use contexts with None and dict."""
    reader = EnvironmentReader(DummyEnv())
    d = {"x": "y"}
    with reader.use(None):
        with reader.use(d):
            pass

def test_use_context_stack_isolation():
    """Test that exiting a context restores stack to previous state."""
    reader = EnvironmentReader(DummyEnv())
    d1 = {"one": 1}
    d2 = {"two": 2}
    # Manually push something
    reader._config_stack.append("manual")
    with reader.use(d1):
        with reader.use(d2):
            pass
    # Cleanup
    reader._config_stack.pop()

def test_use_context_manager_exception():
    """Test that stack is cleaned up even if an exception is raised inside context."""
    reader = EnvironmentReader(DummyEnv())
    d = {"err": True}
    try:
        with reader.use(d):
            raise ValueError("Test error")
    except ValueError:
        pass

def test_use_with_falsey_values():
    """Test that use with falsey values (0, '', False) pushes them as-is."""
    reader = EnvironmentReader(DummyEnv())
    for val in [0, '', False]:
        with reader.use(val):
            pass

def test_use_with_multiple_types():
    """Test that use can handle a variety of types."""
    reader = EnvironmentReader(DummyEnv())
    types = [123, "abc", [1,2], (3,4), set([5]), None, {"x": 99}]
    for val in types:
        with reader.use(val):
            expected = {} if val is None else val

# --------- LARGE SCALE TEST CASES ---------


def test_use_performance_large_data():
    """Test that use can handle large dicts efficiently."""
    reader = EnvironmentReader(DummyEnv())
    large_dict = {str(i): i for i in range(1000)}  # 1000 keys
    with reader.use(large_dict):
        pass

def test_use_large_list_object():
    """Test that use can push a large list object."""
    reader = EnvironmentReader(DummyEnv())
    large_list = list(range(1000))
    with reader.use(large_list):
        pass

def test_use_large_string_object():
    """Test that use can push a large string object."""
    reader = EnvironmentReader(DummyEnv())
    large_str = "x" * 1000
    with reader.use(large_str):
        pass

def test_use_multiple_large_contexts():
    """Test that use can handle multiple large objects in nested contexts."""
    reader = EnvironmentReader(DummyEnv())
    large_dict = {str(i): i for i in range(500)}
    large_list = list(range(500))
    with reader.use(large_dict):
        with reader.use(large_list):
            pass
# 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-EnvironmentReader.use-mgltwy8m and push.