⚡️ Speed up method EnvironmentReader.float by 30%

[codeflash-ai]

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

⏱️ Runtime : 746 microseconds → 576 microseconds (best of 77 runs)

📝 Explanation and details

The optimized code achieves a 29% speedup through several targeted micro-optimizations:

Key optimizations:

1. Faster type checking in read_key(): Replaced isinstance(value, str) with type(value) is str. This avoids the overhead of isinstance's subclass checking since we only care about exact string types, saving ~15% time in this hot function.

2. Tuple instead of list creation: In _read_env(), when converting a single string to an iterable, use (env_key,) tuple instead of [env_key] list. Tuples are faster to create and iterate over for small collections.

3. Eliminated lambda overhead: Replaced lambda k, dv: self._env.float(k, dv) with direct method reference env.float. This removes function call overhead and closure creation, which is significant since _read_env() calls this function in a loop.

4. Reduced attribute access: Cache self._env as local variable env and use getattr(self, "section", None) to safely access the section attribute once instead of repeated property lookups.

Performance characteristics by test case:

• Large-scale tests with many keys see the biggest gains (35-40% faster) due to the eliminated lambda overhead in the tight loop
• Basic single-key lookups see modest 2-7% improvements from the type checking and attribute caching optimizations
• Tests with section overrides benefit from the reduced attribute access patterns

The optimizations are most effective for workloads with large environment key lists or frequent configuration reads, which matches typical usage patterns in configuration management systems.

✅ Correctness verification report:

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

🌀 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 ---

# Helper class to mock Env behavior for testing
class MockEnv:
    def __init__(self, float_map):
        # float_map: dict mapping upper-case keys to float values or string representations
        self.float_map = float_map

    def float(self, key, default):
        # Returns the float value for the key, or default if not found
        val = self.float_map.get(key)
        if val is None:
            return default
        try:
            return float(val)
        except (ValueError, TypeError):
            return default

# Helper Enum for testing
class TestEnum(Enum):
    PI = "PI"
    E = "E"

# 1. Basic Test Cases
def test_float_basic_string_key_found():
    """Test with a string key present in env, value as float."""
    env = MockEnv({"FOO": 3.14})
    reader = EnvironmentReader(env)
    codeflash_output = reader.float("foo") # 3.14μs -> 3.05μs (2.79% faster)

def test_float_basic_string_key_found_str_value():
    """Test with a string key present in env, value as string float."""
    env = MockEnv({"BAR": "2.718"})
    reader = EnvironmentReader(env)
    codeflash_output = reader.float("bar") # 3.09μs -> 3.04μs (1.65% faster)

def test_float_basic_string_key_not_found_returns_default():
    """Test with a string key not present, should return default."""
    env = MockEnv({})
    reader = EnvironmentReader(env)
    codeflash_output = reader.float("baz", default_value=1.23) # 2.87μs -> 2.75μs (3.99% faster)

def test_float_basic_enum_key_found():
    """Test with Enum key present in env."""
    env = MockEnv({"PI": 3.14159})
    reader = EnvironmentReader(env)
    codeflash_output = reader.float(TestEnum.PI) # 3.40μs -> 3.27μs (3.95% faster)

def test_float_basic_enum_key_not_found_returns_default():
    """Test with Enum key not present, should return default."""
    env = MockEnv({})
    reader = EnvironmentReader(env)
    codeflash_output = reader.float(TestEnum.E, default_value=2.71) # 3.32μs -> 3.32μs (0.000% faster)

def test_float_basic_env_key_list_priority():
    """Test with env_key as a list, should prefer first found value."""
    env = MockEnv({"A": 1.1, "B": 2.2})
    reader = EnvironmentReader(env)
    # Only B is present
    codeflash_output = reader.float("foo", env_key=["c", "b"], default_value=0.0) # 3.28μs -> 3.12μs (5.23% faster)
    # Both present, A should be preferred
    env = MockEnv({"A": 1.1, "B": 2.2})
    reader = EnvironmentReader(env)
    codeflash_output = reader.float("foo", env_key=["a", "b"], default_value=0.0) # 1.30μs -> 1.27μs (2.76% faster)



def test_float_edge_env_key_empty_list_returns_default():
    """Test with env_key as empty list, should return default."""
    env = MockEnv({"A": 1.1})
    reader = EnvironmentReader(env)
    codeflash_output = reader.float("foo", env_key=[], default_value=4.2) # 4.11μs -> 3.81μs (7.95% faster)

def test_float_edge_env_key_none_and_key_not_found_returns_default():
    """Test with env_key None and key not found, should return default."""
    env = MockEnv({})
    reader = EnvironmentReader(env)
    codeflash_output = reader.float("missing", env_key=None, default_value=5.5) # 2.94μs -> 2.91μs (0.893% faster)

def test_float_edge_env_key_case_insensitive():
    """Test that key lookup is case-insensitive."""
    env = MockEnv({"FOO": 1.23})
    reader = EnvironmentReader(env)
    codeflash_output = reader.float("Foo") # 2.71μs -> 2.73μs (0.770% slower)
    codeflash_output = reader.float("foo") # 1.11μs -> 1.04μs (6.65% faster)
    codeflash_output = reader.float("FOO") # 736ns -> 700ns (5.14% faster)


def test_float_edge_env_value_not_float_returns_default():
    """Test with env value not convertible to float, should return default."""
    env = MockEnv({"BAD": "not_a_float"})
    reader = EnvironmentReader(env)
    codeflash_output = reader.float("bad", default_value=9.99) # 5.55μs -> 5.36μs (3.73% faster)


def test_float_edge_env_value_none_returns_default():
    """Test with env value as None, should return default."""
    env = MockEnv({"NULL": None})
    reader = EnvironmentReader(env)
    codeflash_output = reader.float("null", default_value=0.01) # 4.11μs -> 3.90μs (5.43% faster)

def test_float_edge_default_none_returns_none():
    """Test with no value found and default_value=None, should return None."""
    env = MockEnv({})
    reader = EnvironmentReader(env)
    codeflash_output = reader.float("missing", default_value=None) # 2.85μs -> 2.83μs (0.849% faster)

def test_float_edge_env_key_list_all_missing_returns_default():
    """Test with env_key list, all missing, should return default."""
    env = MockEnv({})
    reader = EnvironmentReader(env)
    codeflash_output = reader.float("foo", env_key=["x", "y", "z"], default_value=3.3) # 3.29μs -> 3.11μs (5.95% faster)

def test_float_edge_env_key_list_first_value_not_float_second_is_float():
    """Test with env_key list, first value not float, second is float."""
    env = MockEnv({"A": "not_a_float", "B": "2.5"})
    reader = EnvironmentReader(env)
    # Should skip A (returns default), then get B
    codeflash_output = reader.float("foo", env_key=["a", "b"], default_value=0.0) # 5.05μs -> 4.86μs (3.89% faster)

# 3. Large Scale Test Cases
def test_float_large_env_key_list_many_entries():
    """Test with env_key list of 1000 entries, only last has value."""
    keys = [f"KEY{i}" for i in range(1000)]
    env_map = {k.upper(): None for k in keys}
    env_map["KEY999"] = "123.456"
    env = MockEnv(env_map)
    reader = EnvironmentReader(env)
    codeflash_output = reader.float("foo", env_key=keys, default_value=0.0) # 154μs -> 113μs (35.8% faster)

def test_float_large_env_key_list_first_has_value():
    """Test with env_key list of 1000 entries, first has value."""
    keys = [f"KEY{i}" for i in range(1000)]
    env_map = {k.upper(): None for k in keys}
    env_map["KEY0"] = "789.012"
    env = MockEnv(env_map)
    reader = EnvironmentReader(env)
    codeflash_output = reader.float("foo", env_key=keys, default_value=0.0) # 3.36μs -> 3.29μs (2.07% faster)



def test_float_large_env_key_list_all_missing_returns_default():
    """Test with env_key list of 1000 entries, none present, returns default."""
    keys = [f"KEY{i}" for i in range(1000)]
    env_map = {k.upper(): None for k in keys}
    env = MockEnv(env_map)
    reader = EnvironmentReader(env)
    codeflash_output = reader.float("foo", env_key=keys, default_value=7.77) # 155μs -> 112μs (38.3% faster)
# 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 environs import Env
from graphrag.config.environment_reader import EnvironmentReader

# --- Unit Tests ---

# Helper classes for mocking
class MockEnv:
    """Mocked Env class for testing."""
    def __init__(self, float_map=None):
        self.float_map = float_map or {}

    def float(self, key, default):
        # Simulate environs.Env.float: returns float value if key exists, else default
        value = self.float_map.get(key)
        if value is not None:
            return float(value)
        return default

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

@pytest.fixture
def basic_env():
    # Basic environment with some float values
    return MockEnv({
        "FOO": "1.23",
        "BAR": "4.56",
        "BAZ": "7.89",
        "ZERO": "0.0",
        "NEG": "-3.14",
        "EXP": "1e3",
        "NAN": "nan",
        "INF": "inf",
        "NEG_INF": "-inf",
        "COLOR_RED": "42.0",
    })

@pytest.fixture
def env_reader(basic_env):
    return EnvironmentReader(basic_env)

# --- Basic Test Cases ---

def test_basic_float_from_env(env_reader):
    """Test reading basic float values from environment."""
    codeflash_output = env_reader.float("FOO") # 4.04μs -> 4.08μs (0.980% slower)
    codeflash_output = env_reader.float("BAR") # 1.25μs -> 1.23μs (1.71% faster)
    codeflash_output = env_reader.float("BAZ") # 861ns -> 864ns (0.347% slower)

def test_basic_float_with_default(env_reader):
    """Test reading float with default value fallback."""
    # Key not present, should return default
    codeflash_output = env_reader.float("NOT_PRESENT", default_value=9.99) # 2.90μs -> 2.86μs (1.54% faster)

def test_basic_float_zero_and_negative(env_reader):
    """Test reading zero and negative float values."""
    codeflash_output = env_reader.float("ZERO") # 3.10μs -> 2.78μs (11.5% faster)
    codeflash_output = env_reader.float("NEG") # 1.38μs -> 1.42μs (2.75% slower)

def test_basic_float_exponential(env_reader):
    """Test reading float in exponential notation."""
    codeflash_output = env_reader.float("EXP") # 3.00μs -> 2.92μs (2.67% faster)

def test_basic_float_nan_inf(env_reader):
    """Test reading NaN and infinity values."""
    import math
    codeflash_output = env_reader.float("NAN"); nan_val = codeflash_output # 3.01μs -> 2.82μs (6.78% faster)
    codeflash_output = env_reader.float("INF"); inf_val = codeflash_output # 1.28μs -> 1.24μs (3.07% faster)
    codeflash_output = env_reader.float("NEG_INF"); neg_inf_val = codeflash_output # 1.08μs -> 1.01μs (6.61% faster)

def test_basic_float_enum_key(env_reader):
    """Test reading float using Enum key."""
    # Key is Enum, should be converted to lower-case string
    # The environment key will be COLOR_RED (upper-case)
    codeflash_output = env_reader.float(Color.RED, env_key="COLOR_RED") # 3.81μs -> 3.56μs (7.17% faster)

def test_basic_float_list_of_env_keys(env_reader):
    """Test reading float from a list of environment keys."""
    # First key missing, second key present
    codeflash_output = env_reader.float("not_found", env_key=["not_found", "FOO"]) # 3.48μs -> 3.23μs (7.75% faster)

def test_basic_float_section_overrides_env(env_reader):
    """Test that section dictionary overrides environment."""
    env_reader.section = {"foo": "99.99"}
    codeflash_output = env_reader.float("FOO")
    # Section key not present, falls back to env
    env_reader.section = {"bar": "123.45"}
    codeflash_output = env_reader.float("FOO")

# --- Edge Test Cases ---

def test_edge_float_missing_key_no_default(env_reader):
    """Test missing key with no default returns None."""
    codeflash_output = env_reader.float("NOT_PRESENT") # 3.81μs -> 3.70μs (2.89% faster)

def test_edge_float_section_none(env_reader):
    """Test with section set to None."""
    env_reader.section = None
    codeflash_output = env_reader.float("FOO")

def test_edge_float_section_empty(env_reader):
    """Test with section set to empty dict."""
    env_reader.section = {}
    codeflash_output = env_reader.float("FOO")

def test_edge_float_section_case_insensitive(env_reader):
    """Test section key matching is case-insensitive (via lower())."""
    env_reader.section = {"foo": "11.11"}  # lower-case key
    codeflash_output = env_reader.float("FOO")
    env_reader.section = {"FOO": "22.22"}  # upper-case key, but float uses lower-case
    codeflash_output = env_reader.float("foo")

def test_edge_float_env_key_list_all_missing(env_reader):
    """Test env_key list with all missing keys returns default."""
    codeflash_output = env_reader.float("foo", env_key=["not_found1", "not_found2"], default_value=0.5) # 4.27μs -> 4.08μs (4.53% faster)

def test_edge_float_env_key_is_none(env_reader):
    """Test env_key is None, should use key as env_key."""
    codeflash_output = env_reader.float("BAR", env_key=None) # 3.54μs -> 3.48μs (1.64% faster)

def test_edge_float_env_key_is_list_of_one(env_reader):
    """Test env_key is a list with one valid key."""
    codeflash_output = env_reader.float("foo", env_key=["BAR"]) # 3.20μs -> 3.08μs (3.79% faster)

def test_edge_float_env_key_is_empty_list(env_reader):
    """Test env_key is an empty list, should return default."""
    codeflash_output = env_reader.float("foo", env_key=[], default_value=2.2) # 3.24μs -> 3.16μs (2.72% faster)



def test_edge_float_env_key_is_list_with_duplicate_keys(env_reader):
    """Test env_key list with duplicate keys."""
    codeflash_output = env_reader.float("foo", env_key=["FOO", "FOO"]) # 4.31μs -> 4.42μs (2.60% slower)

def test_edge_float_env_key_is_list_with_case_variations(env_reader):
    """Test env_key list with case variations."""
    codeflash_output = env_reader.float("foo", env_key=["foo", "FOO", "Foo"]) # 3.31μs -> 3.23μs (2.63% faster)

def test_edge_float_section_value_is_non_float(env_reader):
    """Test section value is a string representing a float."""
    env_reader.section = {"foo": "3.14159"}
    codeflash_output = env_reader.float("FOO")

def test_edge_float_section_value_is_int_string(env_reader):
    """Test section value is a string representing an int."""
    env_reader.section = {"foo": "42"}
    codeflash_output = env_reader.float("FOO")

def test_edge_float_section_value_is_float_object(env_reader):
    """Test section value is a float object."""
    env_reader.section = {"foo": 2.71828}
    codeflash_output = env_reader.float("FOO")

def test_edge_float_section_value_is_int_object(env_reader):
    """Test section value is an int object."""
    env_reader.section = {"foo": 7}
    codeflash_output = env_reader.float("FOO")

def test_edge_float_section_value_is_invalid(env_reader):
    """Test section value is not convertible to float."""
    env_reader.section = {"foo": "not_a_float"}
    with pytest.raises(ValueError):
        env_reader.float("FOO")

def test_edge_float_env_value_is_invalid(env_reader):
    """Test env value is not convertible to float."""
    env_reader._env.float_map["BAD"] = "not_a_float"
    # Should raise ValueError when trying to convert
    with pytest.raises(ValueError):
        env_reader.float("BAD") # 5.08μs -> 5.05μs (0.614% faster)

def test_edge_float_env_value_is_empty_string(env_reader):
    """Test env value is empty string, should raise ValueError."""
    env_reader._env.float_map["EMPTY"] = ""
    with pytest.raises(ValueError):
        env_reader.float("EMPTY") # 3.71μs -> 3.73μs (0.589% slower)

# --- Large Scale Test Cases ---

def test_large_scale_env_reader_many_keys():
    """Test performance and correctness with many keys in env."""
    # Create env with 1000 float keys
    float_map = {f"KEY{i}": str(i * 0.5) for i in range(1000)}
    env = MockEnv(float_map)
    reader = EnvironmentReader(env)
    # Test a few random keys
    codeflash_output = reader.float("KEY0") # 3.47μs -> 3.26μs (6.63% faster)
    codeflash_output = reader.float("KEY999") # 1.77μs -> 1.73μs (2.49% faster)
    codeflash_output = reader.float("KEY500") # 981ns -> 1.01μs (3.06% slower)
    # Test missing key with default
    codeflash_output = reader.float("NOT_FOUND", default_value=123.456) # 1.56μs -> 1.52μs (2.50% faster)

def test_large_scale_env_reader_many_env_keys():
    """Test with a large env_key list."""
    float_map = {"TARGET": "777.777"}
    env = MockEnv(float_map)
    reader = EnvironmentReader(env)
    # env_key list of 999 missing keys, last one is present
    env_keys = [f"MISS{i}" for i in range(999)] + ["TARGET"]
    codeflash_output = reader.float("foo", env_key=env_keys, default_value=0.0) # 143μs -> 102μs (39.8% faster)



def test_large_scale_env_reader_nan_inf_values():
    """Test large env with many NaN/Inf values."""
    float_map = {f"NAN{i}": "nan" for i in range(500)}
    float_map.update({f"INF{i}": "inf" for i in range(500)})
    env = MockEnv(float_map)
    reader = EnvironmentReader(env)
    import math
    for i in range(0, 500, 100):
        pass

def test_large_scale_env_reader_env_key_case_mismatch():
    """Test large env_key list with case mismatches."""
    float_map = {"SOMEKEY": "123.456"}
    env = MockEnv(float_map)
    reader = EnvironmentReader(env)
    env_keys = ["somekey", "SomeKey", "SOMEKEY"]
    codeflash_output = reader.float("foo", env_key=env_keys, default_value=0.0) # 3.64μs -> 3.64μs (0.055% slower)

def test_large_scale_env_reader_missing_keys():
    """Test large env with all missing keys returns default."""
    float_map = {f"KEY{i}": str(i) for i in range(1000)}
    env = MockEnv(float_map)
    reader = EnvironmentReader(env)
    # All env_keys missing
    env_keys = [f"MISS{i}" for i in range(1000)]
    codeflash_output = reader.float("foo", env_key=env_keys, default_value=1.23) # 152μs -> 111μs (36.6% 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.float-mgluhld0 and push.