⚡️ Speed up function _compute_datetime_types by 10%

[codeflash-ai]

📄 10% (0.10x) speedup for _compute_datetime_types in src/bokeh/util/serialization.py

⏱️ Runtime : 172 microseconds → 157 microseconds (best of 243 runs)

📝 Explanation and details

The optimization achieves a 9% speedup by making two key changes to the _compute_datetime_types() function:

1. Set Literal Construction Optimization:
The original code creates an empty set and uses multiple .add() calls to populate it. The optimized version constructs the set directly using a set literal with all types included upfront. This eliminates the overhead of multiple method calls and intermediate set resizing operations.

2. Import Statement Reorganization:
While the pandas import remains local to the function (preserving lazy loading behavior), the other imports (datetime, numpy) are moved to module level. This reduces the function's execution overhead slightly, though the primary benefit comes from the set construction change.

Why This Works:

• Set literals are faster than incremental construction because Python can allocate the correct size immediately and avoid rehashing
• Eliminating multiple .add() method calls reduces function call overhead
• The @lru_cache(None) decorator ensures this optimization only needs to run once per process in production

Test Case Performance:
The optimization shows consistent 10-20% improvements across most test cases, with particularly strong gains (up to 54%) in cache-miss scenarios. The optimization is most effective for:

• First-time calls to the function (cache misses)
• Applications that clear the cache and recompute
• Any scenario where the function executes its full logic rather than returning cached results

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

import datetime as dt
# function to test
from functools import lru_cache

import numpy as np
# imports
import pytest  # used for our unit tests
from bokeh.util.serialization import _compute_datetime_types

# unit tests

# --- Basic Test Cases ---

def test_contains_datetime_types_basic():
    # Test that the set contains Python's datetime.datetime
    codeflash_output = _compute_datetime_types(); types = codeflash_output # 1.71μs -> 1.41μs (20.7% faster)

def test_contains_time_type_basic():
    # Test that the set contains Python's datetime.time
    codeflash_output = _compute_datetime_types(); types = codeflash_output # 1.49μs -> 1.50μs (0.934% slower)

def test_contains_numpy_datetime64_basic():
    # Test that the set contains numpy.datetime64
    codeflash_output = _compute_datetime_types(); types = codeflash_output # 1.54μs -> 1.38μs (12.1% faster)

def test_contains_pandas_timestamp_basic():
    # Test that the set contains pandas.Timestamp
    import pandas as pd
    codeflash_output = _compute_datetime_types(); types = codeflash_output # 1.43μs -> 1.26μs (13.9% faster)

def test_contains_pandas_timedelta_basic():
    # Test that the set contains pandas.Timedelta
    import pandas as pd
    codeflash_output = _compute_datetime_types(); types = codeflash_output # 1.44μs -> 1.18μs (22.3% faster)

def test_contains_pandas_period_basic():
    # Test that the set contains pandas.Period
    import pandas as pd
    codeflash_output = _compute_datetime_types(); types = codeflash_output # 1.42μs -> 1.22μs (16.8% faster)

def test_contains_pandas_nat_type_basic():
    # Test that the set contains type(pd.NaT)
    import pandas as pd
    codeflash_output = _compute_datetime_types(); types = codeflash_output # 1.45μs -> 1.27μs (13.9% faster)

def test_return_type_is_set():
    # Test that the return type is a set
    codeflash_output = _compute_datetime_types(); types = codeflash_output # 1.51μs -> 1.30μs (15.8% faster)

def test_set_contains_only_types():
    # Test that all elements in the set are types
    codeflash_output = _compute_datetime_types(); types = codeflash_output # 1.54μs -> 1.41μs (9.31% faster)
    for t in types:
        pass

# --- Edge Test Cases ---

def test_no_duplicate_types():
    # Test that the set does not contain duplicate types
    codeflash_output = _compute_datetime_types(); types = codeflash_output # 1.38μs -> 1.38μs (0.073% faster)

def test_set_is_immutable():
    # Test that the returned set is not accidentally mutated
    codeflash_output = _compute_datetime_types(); types = codeflash_output # 1.53μs -> 1.34μs (14.3% faster)
    before = set(types)
    types.add(str)  # Add a new type
    codeflash_output = _compute_datetime_types(); after = codeflash_output # 232ns -> 224ns (3.57% faster)

def test_lru_cache_behavior():
    # Test that lru_cache is working and does not recompute unnecessarily
    import pandas as pd
    _compute_datetime_types.cache_clear()
    codeflash_output = _compute_datetime_types(); types_first = codeflash_output # 1.42μs -> 1.25μs (13.9% faster)
    codeflash_output = _compute_datetime_types(); types_second = codeflash_output # 189ns -> 194ns (2.58% slower)

def test_nat_type_is_not_nat_instance():
    # Test that type(pd.NaT) is not pd.NaT itself
    import pandas as pd
    codeflash_output = _compute_datetime_types(); types = codeflash_output # 1.37μs -> 1.19μs (15.9% faster)

def test_no_unexpected_types():
    # Test that no unexpected types are present
    import pandas as pd
    expected_types = {
        dt.time,
        dt.datetime,
        np.datetime64,
        pd.Timestamp,
        pd.Timedelta,
        pd.Period,
        type(pd.NaT)
    }
    codeflash_output = _compute_datetime_types(); types = codeflash_output # 1.36μs -> 972ns (39.6% faster)

def test_set_is_hashable():
    # Test that the set can be hashed (should fail, sets are not hashable)
    codeflash_output = _compute_datetime_types(); types = codeflash_output # 1.48μs -> 1.36μs (8.66% faster)
    with pytest.raises(TypeError):
        hash(types)

def test_set_is_iterable():
    # Test that the set is iterable
    codeflash_output = _compute_datetime_types(); types = codeflash_output # 1.44μs -> 1.39μs (3.89% faster)
    for t in types:
        pass

# --- Large Scale Test Cases ---

def test_large_scale_membership():
    # Test membership for a large number of objects
    import pandas as pd
    codeflash_output = _compute_datetime_types(); types = codeflash_output # 1.43μs -> 1.29μs (10.7% faster)
    # Create 1000 datetime objects and check their types
    for i in range(1000):
        obj = dt.datetime(2000, 1, 1) if i % 3 == 0 else dt.time(12, 0) if i % 3 == 1 else np.datetime64('2000-01-01')


def test_large_scale_nat_type():
    # Test with many instances of pd.NaT
    import pandas as pd
    codeflash_output = _compute_datetime_types(); types = codeflash_output # 2.57μs -> 2.33μs (10.2% faster)
    for _ in range(1000):
        pass

def test_large_scale_timedelta():
    # Test with many pandas.Timedelta objects
    import pandas as pd
    codeflash_output = _compute_datetime_types(); types = codeflash_output # 1.82μs -> 1.59μs (14.4% faster)
    for i in range(1000):
        td = pd.Timedelta(days=i)

def test_large_scale_timestamp():
    # Test with many pandas.Timestamp objects
    import pandas as pd
    codeflash_output = _compute_datetime_types(); types = codeflash_output # 1.76μs -> 1.55μs (13.6% faster)
    for i in range(1000):
        ts = pd.Timestamp(year=2000, month=1, day=1) + pd.Timedelta(days=i)

def test_large_scale_set_operations():
    # Test set operations with the returned set
    codeflash_output = _compute_datetime_types(); types = codeflash_output # 2.02μs -> 1.77μs (14.1% faster)
    extra_types = {str, int, float}
    union_set = types | extra_types

def test_performance_large_scale():
    # Test that large scale access is efficient (not a real perf test, just that it works)
    codeflash_output = _compute_datetime_types(); types = codeflash_output # 1.78μs -> 1.54μs (15.5% faster)
    # Membership checks for 1000 elements
    for t in list(types) * 100:
        pass
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
import datetime as dt
# function to test
from functools import lru_cache
from typing import Any

import numpy as np
# imports
import pytest  # used for our unit tests
from bokeh.util.serialization import _compute_datetime_types


def __getattr__(name: str) -> Any:
    if name == "DATETIME_TYPES":
        return _compute_datetime_types()
    raise AttributeError

# unit tests

def test_basic_datetime_types_inclusion():
    """Basic: Ensure all expected types are present in the output set."""
    codeflash_output = _compute_datetime_types(); types = codeflash_output # 3.00μs -> 2.70μs (11.0% faster)
    import pandas as pd

def test_basic_return_type_and_uniqueness():
    """Basic: Ensure the function returns a set of types and no duplicates."""
    codeflash_output = _compute_datetime_types(); types = codeflash_output # 2.13μs -> 1.79μs (18.9% faster)
    # All elements should be types
    for t in types:
        pass

def test_edge_type_identity_and_hashability():
    """Edge: Ensure all types in the set are hashable and distinct."""
    codeflash_output = _compute_datetime_types(); types = codeflash_output # 1.89μs -> 1.55μs (21.7% faster)
    # All types should be hashable
    for t in types:
        hash(t)
    # No type should be a subclass of another in the set (except object)
    for t1 in types:
        for t2 in types:
            if t1 is not t2:
                pass

def test_edge_idempotency_and_caching():
    """Edge: Ensure repeated calls return the same object (due to lru_cache)."""
    codeflash_output = _compute_datetime_types(); types1 = codeflash_output # 1.76μs -> 1.41μs (24.7% faster)
    codeflash_output = _compute_datetime_types(); types2 = codeflash_output # 177ns -> 188ns (5.85% slower)

def test_edge_mutability_of_returned_set():
    """Edge: Ensure mutating the returned set does not affect future calls."""
    codeflash_output = _compute_datetime_types(); types1 = codeflash_output # 1.65μs -> 1.54μs (7.02% faster)
    types1.remove(dt.time)
    codeflash_output = _compute_datetime_types(); types2 = codeflash_output # 174ns -> 158ns (10.1% faster)
    # Reset lru_cache for other tests
    _compute_datetime_types.cache_clear()

def test_edge_getattr_success_and_failure():
    """Edge: Test __getattr__ for correct and incorrect attribute access."""
    codeflash_output = _compute_datetime_types() # 1.63μs -> 1.41μs (15.6% faster)
    with pytest.raises(AttributeError):
        __getattr__("NOT_A_REAL_ATTRIBUTE")

def test_edge_type_of_nat_is_not_nat():
    """Edge: Ensure type(pd.NaT) is not pd.NaT itself."""
    import pandas as pd
    codeflash_output = _compute_datetime_types(); types = codeflash_output # 1.51μs -> 1.28μs (17.6% faster)

def test_edge_types_are_not_instances():
    """Edge: Ensure no instance of any type is present in the set."""
    codeflash_output = _compute_datetime_types(); types = codeflash_output # 1.55μs -> 1.43μs (8.62% faster)
    import pandas as pd
    for t in types:
        pass

def test_large_scale_set_usage():
    """Large scale: Test set operations with a large number of elements."""
    codeflash_output = _compute_datetime_types(); types = codeflash_output # 1.59μs -> 1.40μs (13.3% faster)
    # Create a set of 1000 instances of datetime.datetime, all should be recognized by type
    instances = [dt.datetime(2020, 1, 1) for _ in range(1000)]
    # All instances should have type in the types set
    for i in range(1000):
        pass
    # Test set union with a large set of unrelated types
    unrelated_types = {str, int, float, bool, list, dict, tuple}
    union_set = types | unrelated_types

def test_large_scale_type_membership_performance():
    """Large scale: Test membership performance for many types."""
    codeflash_output = _compute_datetime_types(); types = codeflash_output # 1.63μs -> 1.45μs (12.5% faster)
    # Generate 1000 random types, check membership
    for i in range(1000):
        t = type(i)
        if t in types:
            pass
        else:
            pass

def test_large_scale_multiple_calls():
    """Large scale: Ensure multiple calls do not degrade performance or correctness."""
    for _ in range(1000):
        codeflash_output = _compute_datetime_types(); types = codeflash_output # 98.4μs -> 92.2μs (6.76% faster)

def test_large_scale_set_equality():
    """Large scale: Ensure set equality holds for repeated calls."""
    codeflash_output = _compute_datetime_types(); types1 = codeflash_output # 2.57μs -> 2.29μs (12.5% faster)
    codeflash_output = _compute_datetime_types(); types2 = codeflash_output # 170ns -> 166ns (2.41% faster)
    # After cache clear, should still be equal
    _compute_datetime_types.cache_clear()
    codeflash_output = _compute_datetime_types(); types3 = codeflash_output # 912ns -> 592ns (54.1% faster)

def test_edge_types_are_not_builtin_types():
    """Edge: Ensure builtin types like int, float, str are not in the set."""
    codeflash_output = _compute_datetime_types(); types = codeflash_output # 1.50μs -> 1.37μs (9.19% faster)
    for builtin in (int, float, str, bool, list, dict, tuple, set):
        pass

def test_edge_types_are_distinct():
    """Edge: Ensure all types in the set are distinct."""
    codeflash_output = _compute_datetime_types(); types = codeflash_output # 1.63μs -> 1.45μs (12.2% faster)

def test_edge_types_are_not_none_or_object():
    """Edge: Ensure NoneType and object are not in the set."""
    codeflash_output = _compute_datetime_types(); types = codeflash_output # 1.61μs -> 1.44μs (12.0% faster)

def test_edge_types_are_not_functions():
    """Edge: Ensure function types are not in the set."""
    codeflash_output = _compute_datetime_types(); types = codeflash_output # 1.68μs -> 1.52μs (10.6% faster)

def test_edge_types_are_not_classes():
    """Edge: Ensure user-defined classes are not in the set."""
    class Foo: pass
    codeflash_output = _compute_datetime_types(); types = codeflash_output # 1.72μs -> 1.53μs (12.5% faster)

def test_edge_types_are_not_modules():
    """Edge: Ensure module types are not in the set."""
    import sys
    codeflash_output = _compute_datetime_types(); types = codeflash_output # 1.65μs -> 1.45μs (13.8% faster)

def test_edge_types_are_not_numpy_types_other_than_datetime64():
    """Edge: Ensure only numpy.datetime64 is present from numpy types."""
    codeflash_output = _compute_datetime_types(); types = codeflash_output # 1.63μs -> 1.44μs (13.7% faster)
    for t in (np.int32, np.float64, np.bool_):
        pass

def test_edge_types_are_not_pandas_types_other_than_expected():
    """Edge: Ensure only the expected pandas types are present."""
    import pandas as pd
    codeflash_output = _compute_datetime_types(); types = codeflash_output # 1.47μs -> 1.25μs (17.6% faster)
    expected = {pd.Timestamp, pd.Timedelta, pd.Period, type(pd.NaT)}
    for attr in dir(pd):
        obj = getattr(pd, attr)
        if isinstance(obj, type) and obj not in expected:
            pass

def test_edge_types_are_not_datetime_date():
    """Edge: Ensure datetime.date is not in the set."""
    codeflash_output = _compute_datetime_types(); types = codeflash_output # 1.65μs -> 1.44μs (14.9% faster)

def test_edge_types_are_not_datetime_tzinfo():
    """Edge: Ensure datetime.tzinfo is not in the set."""
    codeflash_output = _compute_datetime_types(); types = codeflash_output # 1.63μs -> 1.37μs (18.8% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
from bokeh.util.serialization import _compute_datetime_types
from datetime import datetime
from datetime import time
from numpy import datetime64
from pandas._libs.tslibs.nattype import NaTType
from pandas._libs.tslibs.period import Period
from pandas._libs.tslibs.timedeltas import Timedelta
from pandas._libs.tslibs.timestamps import Timestamp

def test__compute_datetime_types():
    _compute_datetime_types()

To edit these changes git checkout codeflash/optimize-_compute_datetime_types-mhb2n3bd and push.