⚡️ Speed up method EnvironmentReader.str by 30%

[codeflash-ai]

📄 30% (0.30x) speedup for EnvironmentReader.str in graphrag/config/environment_reader.py

⏱️ Runtime : 686 microseconds → 526 microseconds (best of 149 runs)

📝 Explanation and details

The optimized code achieves a 30% speedup through three key optimizations:

1. Eliminated Lambda Creation (Main Performance Gain)
The original code created a new lambda function (lambda k, dv: self._env(k, dv)) on every call to the str method. The optimized version passes self._env directly since it's already a callable with the same signature. This eliminates per-call object allocation overhead, which is the primary source of the performance improvement.

2. Optimized String vs List Handling in _read_env
Instead of always converting single strings to lists (env_key = [env_key]), the optimized version handles string keys directly with immediate lookup and return. This avoids unnecessary list creation and iteration for the common single-key case, reducing overhead by ~24% for string lookups.

3. Cached Section Attribute Access
The original code accessed self.section twice in the conditional check. The optimized version uses getattr(self, "section", None) once and stores it locally, eliminating redundant attribute lookups.

Performance Impact by Test Case:

• Large-scale scenarios benefit most: Tests with 1000+ env_keys show 34%+ speedups due to avoiding repeated lambda allocations
• Basic string lookups: 18-25% faster due to direct string handling and eliminated lambda overhead
• Section lookups: 9-15% faster from cached attribute access
• Single key scenarios: Consistent 20%+ improvements across all test patterns

The optimizations are most effective for workloads with frequent environment variable lookups, especially when using large env_key lists or repeated calls to the str method.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

from collections.abc import Callable
from enum import Enum
from typing import TypeVar

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

# -------------------
# Unit Tests
# -------------------

class DummyEnv:
    """
    Dummy environment class to simulate Env.
    Accepts a dictionary mapping uppercase keys to values.
    """
    def __init__(self, mapping):
        self.mapping = mapping

    def __call__(self, key, default):
        # Simulate environment variable lookup (case-sensitive, as in os.environ)
        return self.mapping.get(key, default)


class Color(Enum):
    RED = "RED"
    GREEN = "GREEN"
    BLUE = "BLUE"

# -------------------
# Basic Test Cases
# -------------------

def test_str_basic_key_present():
    # Test: Key present in env, should return the value
    env = DummyEnv({'FOO': 'bar'})
    reader = EnvironmentReader(env)
    codeflash_output = reader.str('foo') # 3.58μs -> 3.02μs (18.3% faster)

def test_str_basic_key_absent_returns_default():
    # Test: Key not present, should return default_value
    env = DummyEnv({})
    reader = EnvironmentReader(env)
    codeflash_output = reader.str('foo', default_value='baz') # 3.06μs -> 2.53μs (20.9% faster)

def test_str_basic_key_absent_no_default():
    # Test: Key not present, no default_value, should return None
    env = DummyEnv({})
    reader = EnvironmentReader(env)
    codeflash_output = reader.str('foo') # 2.66μs -> 2.15μs (23.8% faster)

def test_str_basic_env_key_list_first_match():
    # Test: env_key as list, first key matches
    env = DummyEnv({'BAR': 'val1', 'BAZ': 'val2'})
    reader = EnvironmentReader(env)
    codeflash_output = reader.str('foo', env_key=['bar', 'baz']) # 2.76μs -> 2.54μs (8.49% faster)

def test_str_basic_env_key_list_second_match():
    # Test: env_key as list, second key matches
    env = DummyEnv({'BAZ': 'val2'})
    reader = EnvironmentReader(env)
    codeflash_output = reader.str('foo', env_key=['bar', 'baz']) # 3.13μs -> 2.62μs (19.7% faster)

def test_str_basic_env_key_list_none_match_returns_default():
    # Test: env_key as list, none match, return default
    env = DummyEnv({})
    reader = EnvironmentReader(env)
    codeflash_output = reader.str('foo', env_key=['bar', 'baz'], default_value='default') # 3.01μs -> 2.69μs (12.0% faster)

def test_str_basic_enum_key():
    # Test: Key is an Enum, should use its value (lowercased)
    env = DummyEnv({'RED': 'apple'})
    reader = EnvironmentReader(env)
    codeflash_output = reader.str(Color.RED) # 3.49μs -> 2.96μs (18.0% faster)

def test_str_basic_env_key_overrides_key():
    # Test: env_key overrides key
    env = DummyEnv({'BAR': 'baz'})
    reader = EnvironmentReader(env)
    codeflash_output = reader.str('foo', env_key='bar') # 2.87μs -> 2.35μs (22.3% faster)

# -------------------
# Edge Test Cases
# -------------------

def test_str_edge_case_sensitive_env_lookup():
    # Test: Environment lookup is case-sensitive (should use .upper())
    env = DummyEnv({'FOO': 'bar', 'foo': 'baz'})
    reader = EnvironmentReader(env)
    # Only 'FOO' should be found, not 'foo'
    codeflash_output = reader.str('foo') # 2.54μs -> 2.14μs (18.7% faster)

def test_str_edge_empty_string_key():
    # Test: Empty string as key
    env = DummyEnv({'': 'empty'})
    reader = EnvironmentReader(env)
    codeflash_output = reader.str('') # 2.54μs -> 2.04μs (24.4% faster)

def test_str_edge_empty_env_key_list():
    # Test: Empty env_key list
    env = DummyEnv({'FOO': 'bar'})
    reader = EnvironmentReader(env)
    codeflash_output = reader.str('foo', env_key=[], default_value='def') # 2.93μs -> 2.38μs (23.1% faster)

def test_str_edge_env_key_is_none():
    # Test: env_key is None, should use key
    env = DummyEnv({'FOO': 'bar'})
    reader = EnvironmentReader(env)
    codeflash_output = reader.str('foo', env_key=None) # 2.78μs -> 2.25μs (23.6% faster)

def test_str_edge_config_stack_section_override():
    # Test: Value present in config_stack section, should override env
    env = DummyEnv({'FOO': 'bar'})
    reader = EnvironmentReader(env)
    reader._config_stack.append({'foo': 'section_value'})
    codeflash_output = reader.str('foo') # 1.64μs -> 1.43μs (14.5% faster)

def test_str_edge_config_stack_section_override_case_insensitive():
    # Test: Section keys are compared lowercased
    env = DummyEnv({'FOO': 'bar'})
    reader = EnvironmentReader(env)
    reader._config_stack.append({'FOO': 'section_value'})
    # Should not match because section key is 'FOO', but lookup is 'foo'
    codeflash_output = reader.str('foo') # 2.88μs -> 2.31μs (24.4% faster)

def test_str_edge_config_stack_and_env_key():
    # Test: Section present, env_key provided, section should win
    env = DummyEnv({'BAR': 'baz'})
    reader = EnvironmentReader(env)
    reader._config_stack.append({'foo': 'section_value'})
    codeflash_output = reader.str('foo', env_key='bar') # 1.79μs -> 1.64μs (9.72% faster)

def test_str_edge_enum_key_with_section():
    # Test: Enum key with section override
    env = DummyEnv({'RED': 'apple'})
    reader = EnvironmentReader(env)
    reader._config_stack.append({'red': 'section_apple'})
    codeflash_output = reader.str(Color.RED) # 2.21μs -> 2.15μs (2.60% faster)

def test_str_edge_default_value_is_none():
    # Test: default_value is None, key not present, should return None
    env = DummyEnv({})
    reader = EnvironmentReader(env)
    codeflash_output = reader.str('foo', default_value=None) # 2.98μs -> 2.39μs (24.6% faster)

def test_str_edge_env_returns_none():
    # Test: Env returns None (explicitly), should return None
    env = DummyEnv({'FOO': None})
    reader = EnvironmentReader(env)
    codeflash_output = reader.str('foo', default_value='bar') # 2.85μs -> 2.37μs (20.2% faster)

def test_str_edge_env_value_is_falsey():
    # Test: Env returns a falsey value (empty string, 0, False)
    env = DummyEnv({'FOO': '', 'BAR': 0, 'BAZ': False})
    reader = EnvironmentReader(env)
    codeflash_output = reader.str('foo', default_value='x') # 2.80μs -> 2.29μs (22.1% faster)
    codeflash_output = reader.str('bar', default_value='x') # 1.15μs -> 1.02μs (12.3% faster)
    codeflash_output = reader.str('baz', default_value='x') # 845ns -> 653ns (29.4% faster)

def test_str_edge_env_key_with_mixed_types():
    # Test: env_key as list with mixed types (string, Enum)
    env = DummyEnv({'RED': 'apple', 'FOO': 'bar'})
    reader = EnvironmentReader(env)
    # Should try 'foo' then 'RED'
    codeflash_output = reader.str('baz', env_key=['foo', Color.RED]) # 2.81μs -> 2.56μs (10.1% faster)

# -------------------
# Large Scale Test Cases
# -------------------

def test_str_large_many_env_keys_first_match():
    # Test: Large env_key list, first key matches
    env_keys = [f'key{i}' for i in range(1000)]
    env = DummyEnv({'KEY0': 'found'})
    reader = EnvironmentReader(env)
    codeflash_output = reader.str('foo', env_key=env_keys) # 2.79μs -> 2.54μs (9.72% faster)

def test_str_large_many_env_keys_last_match():
    # Test: Large env_key list, last key matches
    env_keys = [f'key{i}' for i in range(1000)]
    env = DummyEnv({'KEY999': 'found'})
    reader = EnvironmentReader(env)
    codeflash_output = reader.str('foo', env_key=env_keys) # 147μs -> 110μs (34.1% faster)

def test_str_large_many_env_keys_none_match_returns_default():
    # Test: Large env_key list, none match, return default
    env_keys = [f'key{i}' for i in range(1000)]
    env = DummyEnv({})
    reader = EnvironmentReader(env)
    codeflash_output = reader.str('foo', env_key=env_keys, default_value='default') # 144μs -> 107μs (34.6% faster)

def test_str_large_section_lookup():
    # Test: Large section dictionary, key present
    section = {f'key{i}': f'value{i}' for i in range(1000)}
    env = DummyEnv({'KEY500': 'env_value'})
    reader = EnvironmentReader(env)
    reader._config_stack.append(section)
    codeflash_output = reader.str('key500') # 1.80μs -> 1.64μs (9.77% faster)

def test_str_large_section_lookup_key_absent():
    # Test: Large section dictionary, key absent, should fall back to env
    section = {f'key{i}': f'value{i}' for i in range(1000)}
    env = DummyEnv({'FOO': 'env_value'})
    reader = EnvironmentReader(env)
    reader._config_stack.append(section)
    codeflash_output = reader.str('foo') # 3.10μs -> 2.52μs (22.8% faster)

def test_str_large_env_value_is_large_string():
    # Test: Env value is a large string
    large_value = 'x' * 10000
    env = DummyEnv({'FOO': large_value})
    reader = EnvironmentReader(env)
    codeflash_output = reader.str('foo') # 2.82μs -> 2.28μs (23.6% faster)

def test_str_large_env_key_list_with_enum():
    # Test: Large env_key list with Enum at the end
    env_keys = [f'key{i}' for i in range(999)] + [Color.RED]
    env = DummyEnv({'RED': 'apple'})
    reader = EnvironmentReader(env)
    codeflash_output = reader.str('foo', env_key=env_keys)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
from collections.abc import Callable
from enum import Enum
from typing import TypeVar

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


class DummyEnv:
    """A dummy environment class for testing."""
    def __init__(self, values):
        self.values = {k.upper(): v for k, v in values.items()}

    def __call__(self, key, default):
        # Simulate environment variable lookup (case-insensitive)
        return self.values.get(key.upper(), default)

# unit tests

class Color(Enum):
    RED = "RED"
    BLUE = "BLUE"
    GREEN = "GREEN"

# ---------- Basic Test Cases ----------

def test_str_basic_string_key_found():
    """Test with a basic string key present in env."""
    env = DummyEnv({"foo": "bar"})
    reader = EnvironmentReader(env)
    codeflash_output = reader.str("foo") # 3.62μs -> 3.29μs (10.0% faster)

def test_str_basic_string_key_not_found():
    """Test with a basic string key not present in env, should return None."""
    env = DummyEnv({})
    reader = EnvironmentReader(env)
    codeflash_output = reader.str("foo") # 2.78μs -> 2.25μs (23.6% faster)

def test_str_basic_string_key_with_default():
    """Test with a string key not present, should return default value."""
    env = DummyEnv({})
    reader = EnvironmentReader(env)
    codeflash_output = reader.str("foo", default_value="baz") # 2.92μs -> 2.48μs (17.7% faster)

def test_str_basic_enum_key_found():
    """Test with an Enum key present in env."""
    env = DummyEnv({"red": "apple"})
    reader = EnvironmentReader(env)
    codeflash_output = reader.str(Color.RED) # 3.32μs -> 2.88μs (15.6% faster)

def test_str_basic_enum_key_not_found():
    """Test with an Enum key not present in env, should return None."""
    env = DummyEnv({})
    reader = EnvironmentReader(env)
    codeflash_output = reader.str(Color.BLUE) # 3.23μs -> 2.90μs (11.5% faster)

def test_str_basic_enum_key_with_default():
    """Test with an Enum key not present, should return default value."""
    env = DummyEnv({})
    reader = EnvironmentReader(env)
    codeflash_output = reader.str(Color.GREEN, default_value="leaf") # 3.48μs -> 3.11μs (12.0% faster)

def test_str_basic_env_key_override():
    """Test with env_key overriding key lookup."""
    env = DummyEnv({"bar": "baz"})
    reader = EnvironmentReader(env)
    codeflash_output = reader.str("foo", env_key="bar") # 2.86μs -> 2.34μs (22.2% faster)

def test_str_basic_env_key_list():
    """Test with env_key as a list, first found is returned."""
    env = DummyEnv({"bar": "baz", "qux": "quux"})
    reader = EnvironmentReader(env)
    codeflash_output = reader.str("foo", env_key=["notfound", "qux", "bar"]) # 3.20μs -> 2.90μs (10.3% faster)

def test_str_basic_env_key_list_none_found():
    """Test with env_key as a list, none found, returns default."""
    env = DummyEnv({})
    reader = EnvironmentReader(env)
    codeflash_output = reader.str("foo", env_key=["a", "b"], default_value="def") # 3.00μs -> 2.67μs (12.6% faster)

# ---------- Edge Test Cases ----------

def test_str_edge_key_case_insensitivity():
    """Test that key lookup is case-insensitive."""
    env = DummyEnv({"FOO": "bar"})
    reader = EnvironmentReader(env)
    codeflash_output = reader.str("foo") # 2.59μs -> 2.09μs (23.9% faster)
    codeflash_output = reader.str("FOO") # 1.15μs -> 941ns (22.2% faster)
    codeflash_output = reader.str("FoO") # 794ns -> 608ns (30.6% faster)

def test_str_edge_enum_key_case_insensitivity():
    """Test that Enum key lookup is case-insensitive."""
    env = DummyEnv({"red": "apple", "RED": "cherry"})
    reader = EnvironmentReader(env)
    # Should prefer "red" over "RED" if both exist, but since env keys are uppercased, "RED" wins
    codeflash_output = reader.str(Color.RED) # 3.04μs -> 2.60μs (17.1% faster)

def test_str_edge_empty_env_key_list():
    """Test with empty env_key list, should return default."""
    env = DummyEnv({"foo": "bar"})
    reader = EnvironmentReader(env)
    codeflash_output = reader.str("foo", env_key=[], default_value="baz") # 2.80μs -> 2.28μs (22.9% faster)

def test_str_edge_empty_string_key():
    """Test with empty string key."""
    env = DummyEnv({"": "empty"})
    reader = EnvironmentReader(env)
    codeflash_output = reader.str("") # 2.42μs -> 1.94μs (24.8% faster)

def test_str_edge_special_characters_key():
    """Test with special characters in key."""
    env = DummyEnv({"sp@cial!": "value"})
    reader = EnvironmentReader(env)
    codeflash_output = reader.str("sp@cial!") # 2.63μs -> 2.14μs (22.8% faster)

def test_str_edge_numeric_string_key():
    """Test with numeric string key."""
    env = DummyEnv({"123": "number"})
    reader = EnvironmentReader(env)
    codeflash_output = reader.str("123") # 2.48μs -> 2.05μs (21.3% faster)

def test_str_edge_key_with_whitespace():
    """Test with key containing whitespace."""
    env = DummyEnv({" white space ": "found"})
    reader = EnvironmentReader(env)
    codeflash_output = reader.str(" white space ") # 2.63μs -> 2.09μs (25.9% faster)

def test_str_edge_default_value_none():
    """Test with default_value explicitly set to None."""
    env = DummyEnv({})
    reader = EnvironmentReader(env)
    codeflash_output = reader.str("foo", default_value=None) # 2.92μs -> 2.33μs (25.5% faster)





def test_str_edge_env_key_is_none():
    """Test that if env_key is None, uses key."""
    env = DummyEnv({"foo": "bar"})
    reader = EnvironmentReader(env)
    codeflash_output = reader.str("foo", env_key=None) # 4.01μs -> 3.54μs (13.3% faster)

def test_str_edge_env_key_is_empty_string():
    """Test that if env_key is empty string, returns default."""
    env = DummyEnv({})
    reader = EnvironmentReader(env)
    codeflash_output = reader.str("foo", env_key="", default_value="def") # 3.05μs -> 2.55μs (19.7% faster)


def test_str_large_many_env_keys():
    """Test with a large list of env_keys, only one matches."""
    keys = [f"key{i}" for i in range(999)]
    env = DummyEnv({keys[500]: "found"})
    reader = EnvironmentReader(env)
    # Only the 501st key is present
    codeflash_output = reader.str("foo", env_key=keys, default_value="notfound") # 88.9μs -> 66.6μs (33.6% faster)

def test_str_large_env_key_list_none_found():
    """Test with a large list of env_keys, none matches, returns default."""
    keys = [f"key{i}" for i in range(999)]
    env = DummyEnv({})
    reader = EnvironmentReader(env)
    codeflash_output = reader.str("foo", env_key=keys, default_value="notfound") # 166μs -> 123μs (34.2% faster)

def test_str_large_env_dict():
    """Test with a large env dictionary, key present."""
    env_dict = {f"key{i}": f"val{i}" for i in range(999)}
    env = DummyEnv(env_dict)
    reader = EnvironmentReader(env)
    codeflash_output = reader.str("key123") # 3.08μs -> 2.53μs (21.5% faster)



def test_str_large_env_key_list_multiple_found():
    """Test with a large env_key list, multiple matches, first is returned."""
    keys = [f"key{i}" for i in range(999)]
    env = DummyEnv({keys[10]: "first", keys[20]: "second"})
    reader = EnvironmentReader(env)
    codeflash_output = reader.str("foo", env_key=keys, default_value="notfound") # 5.95μs -> 5.22μs (13.9% 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-EnvironmentReader.str-mglu3mc6 and push.