⚡️ Speed up function _get_api_key_from_environment by 31%

[codeflash-ai]

📄 31% (0.31x) speedup for _get_api_key_from_environment in src/cohere/client.py

⏱️ Runtime : 462 microseconds → 353 microseconds (best of 182 runs)

📝 Explanation and details

The optimized version replaces nested os.getenv() calls with explicit early-return logic using os.environ.get(), providing a ~30% speedup.

Key optimizations:

1. Eliminated nested function calls: The original code always calls os.getenv("COHERE_API_KEY") even when CO_API_KEY is found, due to the nested structure. The optimized version uses early return to avoid the second environment lookup when the preferred key exists.

2. Direct dictionary access: os.environ.get() is faster than os.getenv() because it accesses the environment dictionary directly rather than going through an additional function call wrapper.

Performance characteristics from tests:

• Best case (CO_API_KEY exists): 80-160% faster since it avoids the fallback lookup entirely
• Fallback case (only COHERE_API_KEY exists): 15-80% faster due to the more efficient dictionary access
• No keys case: 20-110% faster from avoiding the nested call overhead

The optimization is particularly effective when CO_API_KEY is set (which is the preferred/documented case), as it completely eliminates the second environment variable lookup. Even in fallback scenarios, the direct dictionary access provides measurable improvements.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

import os
# function to test
import typing

# imports
import pytest
from cohere.client import _get_api_key_from_environment

# --- Basic Test Cases ---

def test_returns_none_when_no_env_vars_set():
    # Neither CO_API_KEY nor COHERE_API_KEY is set
    codeflash_output = _get_api_key_from_environment() # 2.59μs -> 1.22μs (112% faster)

def test_returns_co_api_key_when_set(monkeypatch):
    # CO_API_KEY is set, COHERE_API_KEY is not set
    monkeypatch.setenv("CO_API_KEY", "my-secret-key")
    codeflash_output = _get_api_key_from_environment() # 2.43μs -> 1.05μs (131% faster)

def test_returns_cohere_api_key_when_only_cohere_set(monkeypatch):
    # CO_API_KEY is not set, COHERE_API_KEY is set
    monkeypatch.setenv("COHERE_API_KEY", "another-key")
    codeflash_output = _get_api_key_from_environment() # 1.99μs -> 1.11μs (79.4% faster)

def test_prefers_co_api_key_over_cohere_api_key(monkeypatch):
    # Both CO_API_KEY and COHERE_API_KEY are set, should prefer CO_API_KEY
    monkeypatch.setenv("CO_API_KEY", "preferred-key")
    monkeypatch.setenv("COHERE_API_KEY", "backup-key")
    codeflash_output = _get_api_key_from_environment() # 1.88μs -> 1.00μs (87.2% faster)

def test_returns_empty_string_if_env_var_set_to_empty(monkeypatch):
    # CO_API_KEY is set to empty string
    monkeypatch.setenv("CO_API_KEY", "")
    codeflash_output = _get_api_key_from_environment() # 2.36μs -> 1.03μs (129% faster)
    # CO_API_KEY not set, COHERE_API_KEY is set to empty string
    monkeypatch.delenv("CO_API_KEY", raising=False)
    monkeypatch.setenv("COHERE_API_KEY", "")
    codeflash_output = _get_api_key_from_environment() # 1.43μs -> 1.48μs (3.58% slower)

# --- Edge Test Cases ---

def test_returns_value_with_special_characters(monkeypatch):
    # API key contains special characters
    special_key = "!@#$%^&*()_+-=~`"
    monkeypatch.setenv("CO_API_KEY", special_key)
    codeflash_output = _get_api_key_from_environment() # 2.33μs -> 976ns (138% faster)

def test_returns_value_with_whitespace(monkeypatch):
    # API key contains whitespace
    key_with_spaces = "  key with spaces  "
    monkeypatch.setenv("CO_API_KEY", key_with_spaces)
    codeflash_output = _get_api_key_from_environment() # 2.39μs -> 1.02μs (134% faster)

def test_returns_unicode_value(monkeypatch):
    # API key contains Unicode characters
    unicode_key = "ключ-😀-秘密"
    monkeypatch.setenv("CO_API_KEY", unicode_key)
    codeflash_output = _get_api_key_from_environment() # 3.53μs -> 2.03μs (74.0% faster)

def test_case_sensitivity(monkeypatch):
    # Environment variables are case sensitive
    monkeypatch.setenv("co_api_key", "lowercase-key")
    codeflash_output = _get_api_key_from_environment() # 2.15μs -> 1.01μs (113% faster)
    monkeypatch.setenv("CO_API_KEY", "UPPERCASE-KEY")
    codeflash_output = _get_api_key_from_environment() # 1.39μs -> 744ns (86.2% faster)

def test_cohere_api_key_with_co_api_key_empty(monkeypatch):
    # CO_API_KEY is set to empty string, COHERE_API_KEY is set to a value
    monkeypatch.setenv("CO_API_KEY", "")
    monkeypatch.setenv("COHERE_API_KEY", "backup-key")
    # Should return empty string (CO_API_KEY is preferred even if empty)
    codeflash_output = _get_api_key_from_environment() # 1.84μs -> 1.00μs (83.2% faster)

def test_cohere_api_key_with_co_api_key_unset(monkeypatch):
    # CO_API_KEY is unset, COHERE_API_KEY is set to empty string
    monkeypatch.setenv("COHERE_API_KEY", "")
    codeflash_output = _get_api_key_from_environment() # 1.85μs -> 1.05μs (76.1% faster)

def test_env_var_with_long_value(monkeypatch):
    # Very long API key
    long_key = "a" * 1000
    monkeypatch.setenv("CO_API_KEY", long_key)
    codeflash_output = _get_api_key_from_environment() # 2.60μs -> 1.27μs (105% faster)

def test_env_var_with_newlines(monkeypatch):
    # API key contains newlines
    key_with_newlines = "line1\nline2\nline3"
    monkeypatch.setenv("CO_API_KEY", key_with_newlines)
    codeflash_output = _get_api_key_from_environment() # 2.23μs -> 942ns (137% faster)

# --- Large Scale Test Cases ---

def test_many_unrelated_env_vars(monkeypatch):
    # Set a large number of unrelated environment variables
    for i in range(1000):
        monkeypatch.setenv(f"DUMMY_VAR_{i}", f"value_{i}")
    # Only CO_API_KEY matters
    monkeypatch.setenv("CO_API_KEY", "my-key")
    codeflash_output = _get_api_key_from_environment() # 2.51μs -> 1.35μs (86.1% faster)

def test_many_long_unrelated_env_vars(monkeypatch):
    # Set many unrelated environment variables with long values
    long_value = "x" * 500
    for i in range(1000):
        monkeypatch.setenv(f"UNRELATED_{i}", long_value)
    # Set the correct key
    monkeypatch.setenv("COHERE_API_KEY", "cohere-key")
    codeflash_output = _get_api_key_from_environment() # 2.92μs -> 1.88μs (55.3% faster)

def test_env_var_with_max_length(monkeypatch):
    # Simulate an API key at the upper end of typical environment variable size limits (~32767 on Windows)
    max_length_key = "k" * 4096  # 4KB, well below OS limits, but large for test
    monkeypatch.setenv("CO_API_KEY", max_length_key)
    codeflash_output = _get_api_key_from_environment() # 3.34μs -> 1.66μs (101% faster)

def test_switching_env_vars(monkeypatch):
    # Rapidly change the environment variable and check results
    monkeypatch.setenv("CO_API_KEY", "first")
    codeflash_output = _get_api_key_from_environment() # 2.56μs -> 990ns (159% faster)
    monkeypatch.setenv("CO_API_KEY", "second")
    codeflash_output = _get_api_key_from_environment() # 1.44μs -> 731ns (97.5% faster)
    monkeypatch.delenv("CO_API_KEY", raising=False)
    monkeypatch.setenv("COHERE_API_KEY", "third")
    codeflash_output = _get_api_key_from_environment() # 1.36μs -> 1.42μs (3.87% slower)
    monkeypatch.delenv("COHERE_API_KEY", raising=False)
    codeflash_output = _get_api_key_from_environment() # 1.30μs -> 1.24μs (5.10% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
import os

# imports
import pytest
from cohere.client import _get_api_key_from_environment

# unit tests

@pytest.mark.basic
def test_returns_none_when_no_env_vars_set(monkeypatch):
    """Test case: No relevant environment variables set, should return None."""
    monkeypatch.delenv("CO_API_KEY", raising=False)
    monkeypatch.delenv("COHERE_API_KEY", raising=False)
    codeflash_output = _get_api_key_from_environment() # 2.40μs -> 1.97μs (22.3% faster)

@pytest.mark.basic
def test_returns_co_api_key_when_set(monkeypatch):
    """Test case: Only CO_API_KEY is set, should return its value."""
    monkeypatch.setenv("CO_API_KEY", "abc123")
    monkeypatch.delenv("COHERE_API_KEY", raising=False)
    codeflash_output = _get_api_key_from_environment() # 2.17μs -> 1.02μs (112% faster)

@pytest.mark.basic
def test_returns_cohere_api_key_when_only_cohere_is_set(monkeypatch):
    """Test case: Only COHERE_API_KEY is set, should return its value."""
    monkeypatch.delenv("CO_API_KEY", raising=False)
    monkeypatch.setenv("COHERE_API_KEY", "def456")
    codeflash_output = _get_api_key_from_environment() # 2.07μs -> 1.76μs (17.4% faster)

@pytest.mark.basic
def test_prefers_co_api_key_over_cohere_api_key(monkeypatch):
    """Test case: Both environment variables set, should prefer CO_API_KEY."""
    monkeypatch.setenv("CO_API_KEY", "preferred")
    monkeypatch.setenv("COHERE_API_KEY", "fallback")
    codeflash_output = _get_api_key_from_environment() # 1.89μs -> 1.03μs (82.7% faster)

@pytest.mark.edge
def test_co_api_key_empty_string(monkeypatch):
    """Edge case: CO_API_KEY is set to empty string, should return empty string."""
    monkeypatch.setenv("CO_API_KEY", "")
    monkeypatch.setenv("COHERE_API_KEY", "not_used")
    codeflash_output = _get_api_key_from_environment() # 1.82μs -> 1.02μs (79.2% faster)

@pytest.mark.edge
def test_cohere_api_key_empty_string(monkeypatch):
    """Edge case: COHERE_API_KEY is set to empty string, CO_API_KEY unset."""
    monkeypatch.delenv("CO_API_KEY", raising=False)
    monkeypatch.setenv("COHERE_API_KEY", "")
    codeflash_output = _get_api_key_from_environment() # 2.12μs -> 1.75μs (21.7% faster)

@pytest.mark.edge
def test_co_api_key_none(monkeypatch):
    """Edge case: CO_API_KEY is explicitly set to None (not possible via env, but test for robustness)."""
    monkeypatch.delenv("CO_API_KEY", raising=False)
    monkeypatch.delenv("COHERE_API_KEY", raising=False)
    # os.environ cannot have None values, but test that function doesn't crash
    codeflash_output = _get_api_key_from_environment() # 2.15μs -> 1.75μs (22.8% faster)

@pytest.mark.edge
def test_co_api_key_special_characters(monkeypatch):
    """Edge case: CO_API_KEY contains special characters."""
    special_value = "!@#$%^&*()_+-=[]{}|;':,.<>/?`~"
    monkeypatch.setenv("CO_API_KEY", special_value)
    monkeypatch.delenv("COHERE_API_KEY", raising=False)
    codeflash_output = _get_api_key_from_environment() # 2.08μs -> 1.01μs (106% faster)

@pytest.mark.edge
def test_cohere_api_key_unicode(monkeypatch):
    """Edge case: COHERE_API_KEY contains unicode characters."""
    unicode_value = "ключ🔑"
    monkeypatch.delenv("CO_API_KEY", raising=False)
    monkeypatch.setenv("COHERE_API_KEY", unicode_value)
    codeflash_output = _get_api_key_from_environment() # 3.18μs -> 2.71μs (17.3% faster)

@pytest.mark.edge
def test_both_env_vars_empty(monkeypatch):
    """Edge case: Both environment variables are set to empty string."""
    monkeypatch.setenv("CO_API_KEY", "")
    monkeypatch.setenv("COHERE_API_KEY", "")
    codeflash_output = _get_api_key_from_environment() # 1.75μs -> 1.04μs (68.9% faster)

@pytest.mark.edge
def test_co_api_key_long_string(monkeypatch):
    """Edge case: CO_API_KEY is set to a very long string."""
    long_value = "A" * 1000
    monkeypatch.setenv("CO_API_KEY", long_value)
    monkeypatch.delenv("COHERE_API_KEY", raising=False)
    codeflash_output = _get_api_key_from_environment() # 2.36μs -> 1.32μs (79.5% faster)

@pytest.mark.large
def test_large_number_of_unrelated_env_vars(monkeypatch):
    """Large scale: Many unrelated env vars set, only relevant ones matter."""
    for i in range(1000):
        monkeypatch.setenv(f"UNRELATED_ENV_{i}", str(i))
    monkeypatch.setenv("CO_API_KEY", "large_scale_key")
    monkeypatch.setenv("COHERE_API_KEY", "should_not_be_used")
    codeflash_output = _get_api_key_from_environment() # 2.12μs -> 1.35μs (56.9% faster)

@pytest.mark.large
def test_large_cohere_api_key(monkeypatch):
    """Large scale: COHERE_API_KEY is set to a very large string, CO_API_KEY unset."""
    large_value = "X" * 999
    monkeypatch.delenv("CO_API_KEY", raising=False)
    monkeypatch.setenv("COHERE_API_KEY", large_value)
    codeflash_output = _get_api_key_from_environment() # 2.58μs -> 2.24μs (15.1% faster)

@pytest.mark.large
def test_many_calls_with_varied_env(monkeypatch):
    """Large scale: Simulate many calls with different env values to check determinism."""
    for i in range(100):
        monkeypatch.setenv("CO_API_KEY", f"key_{i}")
        monkeypatch.setenv("COHERE_API_KEY", f"fallback_{i}")
        codeflash_output = _get_api_key_from_environment() # 102μs -> 60.1μs (70.0% faster)
        monkeypatch.delenv("CO_API_KEY", raising=False)
        codeflash_output = _get_api_key_from_environment()
        monkeypatch.delenv("COHERE_API_KEY", raising=False) # 113μs -> 105μs (7.53% faster)
        codeflash_output = _get_api_key_from_environment()
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
from cohere.client import _get_api_key_from_environment

def test__get_api_key_from_environment():
    _get_api_key_from_environment()

⏪ Replay Tests and Runtime

Test File::Test Function	Original ⏱️	Optimized ⏱️	Speedup
test_pytest_teststest_embed_utils_py_teststest_client_py_teststest_client_init_py_teststest_async_client___replay_test_0.py::test_cohere_client__get_api_key_from_environment	37.1μs	19.7μs	88.6%✅
test_pytest_teststest_pytest_tests__replay_test_0_py_teststest_bedrock_client_py_teststest_overrides_py_t__replay_test_0.py::test_cohere_client__get_api_key_from_environment	5.04μs	2.53μs	99.4%✅

🔎 Concolic Coverage Tests and Runtime

Test File::Test Function	Original ⏱️	Optimized ⏱️	Speedup
codeflash_concolic_yxtehl4j/tmp5ewogeh3/test_concolic_coverage.py::test__get_api_key_from_environment	2.81μs	1.27μs	122%✅

To edit these changes git checkout codeflash/optimize-_get_api_key_from_environment-mgzrkks2 and push.