⚡️ Speed up function _valid_locales by 16%

[codeflash-ai]

📄 16% (0.16x) speedup for _valid_locales in pandas/_config/localization.py

⏱️ Runtime : 16.8 milliseconds → 14.5 milliseconds (best of 63 runs)

📝 Explanation and details

The optimization achieves a 15% speedup through two key changes:

1. Direct locale manipulation in can_set_locale:

• Replaced the set_locale context manager call with direct locale.setlocale() operations
• Eliminates context manager overhead (generator setup, try/finally blocks, function calls)
• Uses a simple try/finally pattern to save/restore the current locale, avoiding the more complex context manager machinery

2. Explicit loop in _valid_locales:

• Replaced the nested list comprehension with a straightforward for-loop and append()
• Reduces intermediate object creation and memory allocation overhead
• Improves CPU cache locality by processing one locale at a time sequentially

Performance impact by test type:

• Single locale tests: 37-124% faster due to reduced context manager overhead
• Multiple locale tests: 8-95% faster, with larger gains when more valid locales are processed
• Large-scale tests: 9-45% faster, showing consistent improvements even with 1000+ locales
• Edge cases: 15-102% faster for empty inputs and invalid locales

The optimizations are particularly effective for workloads that validate many locales, as each can_set_locale call now has significantly less function call and context management overhead. The changes maintain identical functionality while being more direct and efficient.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

from __future__ import annotations

import locale
from collections.abc import Generator
from contextlib import contextmanager

# imports
import pytest  # used for our unit tests
from pandas._config.localization import _valid_locales


@contextmanager
def set_locale(
    new_locale: str | tuple[str, str], lc_var: int = locale.LC_ALL
) -> Generator[str | tuple[str, str]]:
    current_locale = locale.setlocale(lc_var)
    try:
        locale.setlocale(lc_var, new_locale)
        normalized_code, normalized_encoding = locale.getlocale()
        if normalized_code is not None and normalized_encoding is not None:
            yield f"{normalized_code}.{normalized_encoding}"
        else:
            yield new_locale
    finally:
        locale.setlocale(lc_var, current_locale)
from pandas._config.localization import _valid_locales

# unit tests

# Helper: Find some valid locales on the current system for robust tests
def get_some_valid_locales(n=3):
    # Try to get at least n valid locales from locale.locale_alias
    found = []
    for loc in set(locale.locale_alias.values()):
        # Filter out empty and non-UTF8 locales for reliability
        if "." not in loc or "utf" not in loc.lower():
            continue
        try:
            with set_locale(loc):
                found.append(loc)
        except Exception:
            continue
        if len(found) >= n:
            break
    # If not enough found, fallback to C.UTF-8 and en_US.UTF-8
    if len(found) < n:
        for fallback in ["C.UTF-8", "en_US.UTF-8", "POSIX"]:
            try:
                with set_locale(fallback):
                    found.append(fallback)
            except Exception:
                continue
            if len(found) >= n:
                break
    return found[:n]

# Basic Test Cases

def test_single_valid_locale_no_normalize():
    # Use a locale that should be valid on most systems
    valid_locale = get_some_valid_locales(1)[0]
    codeflash_output = _valid_locales([valid_locale], normalize=False); result = codeflash_output # 14.6μs -> 6.51μs (124% faster)

def test_single_invalid_locale_no_normalize():
    # Use a clearly invalid locale string
    invalid_locale = "xx_XX.INVALID"
    codeflash_output = _valid_locales([invalid_locale], normalize=False); result = codeflash_output # 11.2μs -> 8.11μs (37.7% faster)

def test_multiple_locales_mixed_validity_no_normalize():
    valid_locale = get_some_valid_locales(1)[0]
    invalid_locale = "xx_XX.INVALID"
    codeflash_output = _valid_locales([valid_locale, invalid_locale], normalize=False); result = codeflash_output # 18.3μs -> 9.39μs (95.2% faster)

def test_multiple_locales_all_valid_no_normalize():
    valid_locales = get_some_valid_locales(2)
    codeflash_output = _valid_locales(valid_locales, normalize=False); result = codeflash_output # 35.1μs -> 21.0μs (66.6% faster)

def test_multiple_locales_all_invalid_no_normalize():
    invalid_locales = ["bad_LOCALE", "nope_NO.UTF-8", "12345"]
    codeflash_output = _valid_locales(invalid_locales, normalize=False); result = codeflash_output # 76.8μs -> 70.7μs (8.74% faster)

def test_strip_whitespace_from_locales():
    valid_locale = get_some_valid_locales(1)[0]
    codeflash_output = _valid_locales([f"   {valid_locale}   "], normalize=False); result = codeflash_output # 14.3μs -> 6.38μs (124% faster)

def test_input_as_string_instead_of_list():
    # Should treat string as iterable of characters, so expect only valid single-char locales
    valid_locale = get_some_valid_locales(1)[0]
    codeflash_output = _valid_locales(valid_locale, normalize=False); result = codeflash_output # 156μs -> 141μs (10.8% faster)

def test_empty_list_input():
    codeflash_output = _valid_locales([], normalize=False); result = codeflash_output # 1.37μs -> 677ns (102% faster)

def test_empty_string_input():
    codeflash_output = _valid_locales("", normalize=False); result = codeflash_output # 1.32μs -> 780ns (69.9% faster)

# Edge Test Cases

def test_locale_with_extra_newlines_and_spaces():
    valid_locale = get_some_valid_locales(1)[0]
    input_locales = [f"\n{valid_locale}\n", "   ", "\t"]
    codeflash_output = _valid_locales(input_locales, normalize=False); result = codeflash_output # 57.9μs -> 41.8μs (38.5% faster)

def test_locale_with_weird_characters():
    weird_locale = "en_US.UTF-8!@#$%^&*()"
    codeflash_output = _valid_locales([weird_locale], normalize=False); result = codeflash_output # 32.1μs -> 27.9μs (15.3% faster)

def test_locale_case_sensitivity():
    valid_locale = get_some_valid_locales(1)[0]
    lower = valid_locale.lower()
    upper = valid_locale.upper()
    # At least one of these should be valid if the system is case-insensitive
    codeflash_output = _valid_locales([lower, upper], normalize=False); result = codeflash_output # 27.6μs -> 24.4μs (13.1% faster)


def test_locale_with_tuple_input():
    # Should raise, as tuples are not supported
    with pytest.raises(AttributeError):
        _valid_locales([("en_US", "UTF-8")], normalize=False) # 2.64μs -> 1.72μs (53.2% faster)

def test_locale_with_int_input():
    # Should raise, as int is not a valid locale
    with pytest.raises(AttributeError):
        _valid_locales([123], normalize=False) # 2.53μs -> 1.68μs (51.0% faster)

def test_locale_with_none_input():
    # Should raise, as None is not a valid locale
    with pytest.raises(AttributeError):
        _valid_locales([None], normalize=False) # 2.38μs -> 1.61μs (47.3% faster)

def test_locale_with_only_whitespace():
    codeflash_output = _valid_locales(["   ", "\t", "\n"], normalize=False); result = codeflash_output # 131μs -> 106μs (23.2% faster)


def test_locale_with_normalize_false():
    alias = "en_us.utf8"
    codeflash_output = _valid_locales([alias], normalize=False); result = codeflash_output # 47.8μs -> 42.7μs (12.1% faster)

# Large Scale Test Cases


def test_large_number_of_invalid_locales():
    # Generate 100 invalid locales
    invalid_locales = [f"bad_LOCALE_{i}" for i in range(100)]
    codeflash_output = _valid_locales(invalid_locales, normalize=False); result = codeflash_output # 1.21ms -> 1.10ms (10.2% faster)

def test_large_mixed_locales():
    valid_locales = get_some_valid_locales(5)
    invalid_locales = [f"bad_LOCALE_{i}" for i in range(995)]
    input_locales = valid_locales + invalid_locales
    codeflash_output = _valid_locales(input_locales, normalize=False); result = codeflash_output # 11.3ms -> 10.4ms (9.02% faster)



#------------------------------------------------
from __future__ import annotations

import locale  # used to get available locales and normalization
from collections.abc import Generator
from contextlib import contextmanager

# imports
import pytest  # used for our unit tests
from pandas._config.localization import _valid_locales

# unit tests

# Helper: Get a list of available locales on the system
def get_available_locales():
    # Try to get all locales available on the system
    try:
        available = set(locale.locale_alias.values())
        # Remove duplicates and filter out empty strings
        return [l for l in set(available) if l and isinstance(l, str)]
    except Exception:
        # Fallback: Just use some well-known locales
        return ['en_US.UTF-8', 'C', 'POSIX']

# Helper: Find at least one valid and one invalid locale for testing
def get_one_valid_locale():
    for loc in get_available_locales():
        if can_set_locale(loc):
            return loc
    return 'C'  # 'C' locale is almost always available

def get_one_invalid_locale():
    # Try a locale that is almost certainly invalid
    return 'invalid_LOCALE.UTF-8'

# 1. Basic Test Cases



def test_single_invalid_locale():
    # Test with a single invalid locale
    invalid_locale = get_one_invalid_locale()
    codeflash_output = _valid_locales([invalid_locale], normalize=True); result = codeflash_output # 37.1μs -> 37.0μs (0.232% faster)



def test_empty_list():
    # Test with an empty list
    codeflash_output = _valid_locales([], normalize=True); result = codeflash_output # 1.66μs -> 841ns (97.0% faster)




def test_locale_with_empty_string():
    # Test with an empty string as a locale
    codeflash_output = _valid_locales([''], normalize=True); result = codeflash_output # 94.7μs -> 86.2μs (9.86% faster)

def test_locale_with_only_spaces():
    # Test with a locale that is only spaces
    codeflash_output = _valid_locales(['   '], normalize=True); result = codeflash_output # 28.7μs -> 24.7μs (16.2% faster)

def test_locale_with_special_characters():
    # Test with a locale containing special characters
    codeflash_output = _valid_locales(['@@@###$], normalize=True); result = codeflash_output # 19.1μs -> 15.1μs (26.4% faster)

def test_locale_with_none_in_list():
    # Test with None in the list (should raise AttributeError on strip)
    with pytest.raises(AttributeError):
        _valid_locales([None], normalize=True) # 2.68μs -> 2.05μs (30.6% faster)

def test_locale_with_non_string_type():
    # Test with an integer in the list (should raise AttributeError on strip)
    with pytest.raises(AttributeError):
        _valid_locales([123], normalize=True) # 2.75μs -> 1.87μs (47.4% faster)

def test_locale_with_tuple_in_list():
    # Test with a tuple in the list (should raise AttributeError on strip)
    with pytest.raises(AttributeError):
        _valid_locales([(1,2)], normalize=True) # 2.72μs -> 1.95μs (39.1% faster)





def test_large_number_of_locales_all_invalid():
    # Test with a large list of invalid locales
    invalid_locale = get_one_invalid_locale()
    locales = [invalid_locale] * 1000  # 1000 invalid locales
    codeflash_output = _valid_locales(locales, normalize=True); result = codeflash_output # 3.44ms -> 2.37ms (45.2% faster)

To edit these changes git checkout codeflash/optimize-_valid_locales-mhbmrahj and push.