⚡️ Speed up function _check_storage_url by 5%

[codeflash-ai]

📄 5% (0.05x) speedup for _check_storage_url in optuna/cli.py

⏱️ Runtime : 132 microseconds → 125 microseconds (best of 284 runs)

📝 Explanation and details

The optimized code replaces os.environ.get("OPTUNA_STORAGE") with direct dictionary access os.environ["OPTUNA_STORAGE"] wrapped in a try/except block. This eliminates the method call overhead of .get() and its internal null checks.

Key optimizations:

• Direct dictionary access: os.environ["OPTUNA_STORAGE"] avoids the .get() method call overhead and associated branching logic
• Exception-based flow control: Using try/except KeyError is more efficient than checking if env_storage is not None since KeyError handling is optimized in CPython
• Reduced branching: The original code has two conditional checks (storage_url check + env_storage check), while the optimized version uses exception handling to eliminate one branch

Performance gains by test case:

• Fastest improvement (28.3%): When environment variable is missing and CLIUsageError is raised - the direct KeyError exception path is much faster than the .get() + None check + raise sequence
• Moderate improvements (5-8%): When storage_url is provided directly, avoiding any environment variable lookup entirely
• Consistent gains: Most test cases show 1-8% improvements, with the optimization being particularly effective for error cases where the environment variable is not set

The optimization is most beneficial for CLI applications where storage URLs are frequently validated and environment variables may often be unset.

✅ Correctness verification report:

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

⚙️ Existing Unit Tests and Runtime

Test File::Test Function	Original ⏱️	Optimized ⏱️	Speedup
test_cli.py::test_check_storage_url	9.74μs	8.69μs	12.0%✅

🌀 Generated Regression Tests and Runtime

from __future__ import annotations

import os
import warnings

# imports
import pytest
from optuna.cli import _check_storage_url
from optuna.exceptions import CLIUsageError, ExperimentalWarning

# --------------------------
# 1. Basic Test Cases
# --------------------------

def test_storage_url_given_returns_value():
    # If storage_url is given, it should be returned as-is
    codeflash_output = _check_storage_url("sqlite:///foo.db") # 412ns -> 387ns (6.46% faster)
    codeflash_output = _check_storage_url("mysql://user:pass@host/db") # 196ns -> 223ns (12.1% slower)
    codeflash_output = _check_storage_url("") # 131ns -> 129ns (1.55% faster)

def test_storage_url_given_none_and_env_set(monkeypatch):
    # If storage_url is None and OPTUNA_STORAGE is set, should return env value and warn
    monkeypatch.setenv("OPTUNA_STORAGE", "sqlite:///bar.db")
    with warnings.catch_warnings(record=True) as w:
        warnings.simplefilter("always")
        codeflash_output = _check_storage_url(None); result = codeflash_output # 5.81μs -> 5.70μs (1.95% faster)

def test_storage_url_given_none_and_env_not_set():
    # If storage_url is None and OPTUNA_STORAGE is not set, should raise CLIUsageError
    with pytest.raises(CLIUsageError) as excinfo:
        _check_storage_url(None) # 2.47μs -> 1.93μs (28.3% faster)

# --------------------------
# 2. Edge Test Cases
# --------------------------

def test_storage_url_is_empty_string(monkeypatch):
    # If storage_url is "", it should be returned (not treated as None)
    codeflash_output = _check_storage_url("") # 394ns -> 375ns (5.07% faster)

def test_env_var_is_empty_string(monkeypatch):
    # If OPTUNA_STORAGE is set to empty string, it should be returned and warning emitted
    monkeypatch.setenv("OPTUNA_STORAGE", "")
    with warnings.catch_warnings(record=True) as w:
        warnings.simplefilter("always")
        codeflash_output = _check_storage_url(None); result = codeflash_output # 5.62μs -> 5.59μs (0.501% faster)

def test_storage_url_is_whitespace(monkeypatch):
    # If storage_url is whitespace, it should be returned as-is
    codeflash_output = _check_storage_url("   ") # 353ns -> 338ns (4.44% faster)

def test_env_var_is_whitespace(monkeypatch):
    # If OPTUNA_STORAGE is whitespace, it should be returned and warning emitted
    monkeypatch.setenv("OPTUNA_STORAGE", "   ")
    with warnings.catch_warnings(record=True) as w:
        warnings.simplefilter("always")
        codeflash_output = _check_storage_url(None); result = codeflash_output # 5.36μs -> 5.50μs (2.51% slower)

def test_storage_url_is_special_chars(monkeypatch):
    # Accepts special characters in storage_url
    special_url = "sqlite:///:memory:?cache=shared&mode=memory"
    codeflash_output = _check_storage_url(special_url) # 344ns -> 353ns (2.55% slower)

def test_env_var_is_special_chars(monkeypatch):
    # Accepts special characters in OPTUNA_STORAGE
    special_url = "sqlite:///:memory:?cache=shared&mode=memory"
    monkeypatch.setenv("OPTUNA_STORAGE", special_url)
    with warnings.catch_warnings(record=True) as w:
        warnings.simplefilter("always")
        codeflash_output = _check_storage_url(None); result = codeflash_output # 5.31μs -> 5.41μs (1.92% slower)




def test_env_var_is_long_string(monkeypatch):
    # Test with a very long string in OPTUNA_STORAGE
    long_url = "sqlite:///" + "a" * 500
    monkeypatch.setenv("OPTUNA_STORAGE", long_url)
    with warnings.catch_warnings(record=True) as w:
        warnings.simplefilter("always")
        codeflash_output = _check_storage_url(None); result = codeflash_output # 7.67μs -> 7.60μs (0.960% faster)

# --------------------------
# 3. Large Scale Test Cases
# --------------------------

def test_env_var_many_chars(monkeypatch):
    # Test with a storage URL of 999 characters
    url = "sqlite:///" + "x" * 990
    monkeypatch.setenv("OPTUNA_STORAGE", url)
    with warnings.catch_warnings(record=True) as w:
        warnings.simplefilter("always")
        codeflash_output = _check_storage_url(None); result = codeflash_output # 5.79μs -> 5.54μs (4.55% faster)

def test_storage_url_many_chars(monkeypatch):
    # Test with a direct storage_url of 999 characters
    url = "sqlite:///" + "y" * 990
    codeflash_output = _check_storage_url(url) # 380ns -> 364ns (4.40% faster)

def test_env_var_many_unique_urls(monkeypatch):
    # Test repeatedly with different env values (simulate many unique URLs)
    for i in range(10):  # keep <1000 for performance
        url = f"sqlite:///db_{i}.sqlite"
        monkeypatch.setenv("OPTUNA_STORAGE", url)
        with warnings.catch_warnings(record=True) as w:
            warnings.simplefilter("always")
            codeflash_output = _check_storage_url(None); result = codeflash_output

def test_storage_url_many_unique(monkeypatch):
    # Test repeatedly with different direct storage_url values
    for i in range(10):  # keep <1000 for performance
        url = f"sqlite:///foo_{i}.db"
        codeflash_output = _check_storage_url(url) # 1.54μs -> 1.53μs (0.718% faster)

def test_env_var_unicode(monkeypatch):
    # Test with unicode characters in OPTUNA_STORAGE
    unicode_url = "sqlite:///データベース.db"
    monkeypatch.setenv("OPTUNA_STORAGE", unicode_url)
    with warnings.catch_warnings(record=True) as w:
        warnings.simplefilter("always")
        codeflash_output = _check_storage_url(None); result = codeflash_output # 5.84μs -> 5.65μs (3.45% faster)

def test_storage_url_unicode(monkeypatch):
    # Test with unicode characters in direct storage_url
    unicode_url = "sqlite:///база_данных.db"
    codeflash_output = _check_storage_url(unicode_url) # 332ns -> 338ns (1.78% slower)
# 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 os
import warnings

# imports
import pytest  # used for our unit tests
from optuna.cli import _check_storage_url
from optuna.exceptions import CLIUsageError, ExperimentalWarning

# Basic Test Cases

def test_storage_url_direct_string():
    # Should return the provided storage_url directly
    url = "sqlite:///example.db"
    codeflash_output = _check_storage_url(url) # 402ns -> 372ns (8.06% faster)

def test_storage_url_direct_empty_string():
    # Should return empty string if that's the explicit input
    url = ""
    codeflash_output = _check_storage_url(url) # 340ns -> 331ns (2.72% faster)

def test_storage_url_direct_none_env_set(monkeypatch):
    # Should fallback to env var if storage_url is None
    env_url = "mysql://user:pass@localhost/db"
    monkeypatch.setenv("OPTUNA_STORAGE", env_url)
    with warnings.catch_warnings(record=True) as w:
        warnings.simplefilter("always")
        codeflash_output = _check_storage_url(None); result = codeflash_output # 5.77μs -> 5.46μs (5.70% faster)

def test_storage_url_direct_none_env_not_set():
    # Should raise CLIUsageError if both are None
    with pytest.raises(CLIUsageError):
        _check_storage_url(None) # 2.62μs -> 2.16μs (21.4% faster)

# Edge Test Cases

def test_storage_url_none_env_empty_string(monkeypatch):
    # Should return empty string from env if storage_url is None and env is empty string
    monkeypatch.setenv("OPTUNA_STORAGE", "")
    with warnings.catch_warnings(record=True) as w:
        warnings.simplefilter("always")
        codeflash_output = _check_storage_url(None); result = codeflash_output # 5.45μs -> 5.07μs (7.66% faster)

def test_storage_url_none_env_whitespace(monkeypatch):
    # Should return whitespace string from env if set
    monkeypatch.setenv("OPTUNA_STORAGE", "   ")
    with warnings.catch_warnings(record=True) as w:
        warnings.simplefilter("always")
        codeflash_output = _check_storage_url(None); result = codeflash_output # 4.99μs -> 4.71μs (5.86% faster)

def test_storage_url_env_special_characters(monkeypatch):
    # Should handle special characters in env var
    special_url = "postgresql://usér:päss@host/db?param=✓"
    monkeypatch.setenv("OPTUNA_STORAGE", special_url)
    with warnings.catch_warnings(record=True) as w:
        warnings.simplefilter("always")
        codeflash_output = _check_storage_url(None); result = codeflash_output # 5.93μs -> 5.51μs (7.47% faster)

def test_storage_url_env_long_string(monkeypatch):
    # Should handle a long string in env var
    long_url = "sqlite:///" + "a" * 500
    monkeypatch.setenv("OPTUNA_STORAGE", long_url)
    with warnings.catch_warnings(record=True) as w:
        warnings.simplefilter("always")
        codeflash_output = _check_storage_url(None); result = codeflash_output # 4.98μs -> 4.59μs (8.45% faster)

def test_storage_url_env_unicode(monkeypatch):
    # Should handle unicode characters in env var
    unicode_url = "sqlite:///データベース.db"
    monkeypatch.setenv("OPTUNA_STORAGE", unicode_url)
    with warnings.catch_warnings(record=True) as w:
        warnings.simplefilter("always")
        codeflash_output = _check_storage_url(None); result = codeflash_output # 5.26μs -> 5.04μs (4.34% faster)

# Large Scale Test Cases

def test_storage_url_env_large(monkeypatch):
    # Should handle very large env var (close to 1000 chars)
    large_url = "sqlite:///" + "x" * 990
    monkeypatch.setenv("OPTUNA_STORAGE", large_url)
    with warnings.catch_warnings(record=True) as w:
        warnings.simplefilter("always")
        codeflash_output = _check_storage_url(None); result = codeflash_output # 4.92μs -> 4.61μs (6.72% faster)

def test_storage_url_direct_large():
    # Should handle very large direct input
    large_url = "sqlite:///" + "y" * 990
    codeflash_output = _check_storage_url(large_url) # 393ns -> 370ns (6.22% faster)

def test_storage_url_env_var_overwrites_none(monkeypatch):
    # Should always prefer direct input over env, even if env is set
    monkeypatch.setenv("OPTUNA_STORAGE", "sqlite:///env.db")
    direct_url = "sqlite:///direct.db"
    # Should return direct_url, not env var
    codeflash_output = _check_storage_url(direct_url) # 336ns -> 345ns (2.61% slower)

def test_storage_url_env_var_case_sensitive(monkeypatch):
    # Should only use 'OPTUNA_STORAGE', not lowercased or other variants
    monkeypatch.setenv("optuna_storage", "sqlite:///wrong.db")
    # 'OPTUNA_STORAGE' is not set, so should raise
    with pytest.raises(CLIUsageError):
        _check_storage_url(None) # 1.92μs -> 1.72μs (11.2% faster)

def test_storage_url_env_var_removed(monkeypatch):
    # Should raise if env var is removed after being set
    monkeypatch.setenv("OPTUNA_STORAGE", "sqlite:///temp.db")
    monkeypatch.delenv("OPTUNA_STORAGE")
    with pytest.raises(CLIUsageError):
        _check_storage_url(None) # 1.93μs -> 1.62μs (19.6% faster)

def test_storage_url_env_var_multiple(monkeypatch):
    # Should only use 'OPTUNA_STORAGE', not other storage-related vars
    monkeypatch.setenv("OPTUNA_STORAGE", "sqlite:///main.db")
    monkeypatch.setenv("OTHER_STORAGE", "sqlite:///other.db")
    with warnings.catch_warnings(record=True) as w:
        warnings.simplefilter("always")
        codeflash_output = _check_storage_url(None); result = codeflash_output # 5.89μs -> 6.00μs (1.80% slower)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
from optuna.cli import _check_storage_url
import pytest

def test__check_storage_url():
    with pytest.raises(CLIUsageError, match='Storage\\ URL\\ is\\ not\\ specified\\.'):
        _check_storage_url(None)

def test__check_storage_url_2():
    _check_storage_url('')

🔎 Concolic Coverage Tests and Runtime

Test File::Test Function	Original ⏱️	Optimized ⏱️	Speedup
codeflash_concolic_wou29s7s/tmpwi4mrh73/test_concolic_coverage.py::test__check_storage_url_2	501ns	486ns	3.09%✅

To edit these changes git checkout codeflash/optimize-_check_storage_url-mhb1fmr7 and push.