⚡️ Speed up method WSConnector._get_connect_args by 139%

[codeflash-ai]

📄 139% (1.39x) speedup for WSConnector._get_connect_args in distributed/comm/ws.py

⏱️ Runtime : 45.1 microseconds → 18.9 microseconds (best of 567 runs)

📝 Explanation and details

The optimization replaces dictionary unpacking ({**connection_args.get("extra_conn_args", {})}) with a direct return using the or operator (connection_args.get("extra_conn_args") or {}).

Key changes:

• Eliminates unnecessary dictionary construction when extra_conn_args already exists as a dictionary
• Uses Python's or operator to handle the fallback case instead of providing a default empty dict to .get()

Why it's faster:
The original code always creates a new dictionary through unpacking, even when extra_conn_args is already a valid dictionary. Dictionary unpacking involves iterating through key-value pairs and constructing a new dict object. The optimized version returns the existing dictionary directly when it exists, only creating a new empty dict when needed.

Performance characteristics:
The optimization shows dramatic improvements for large-scale test cases (290-529% faster) because it completely avoids dictionary construction overhead. For small dictionaries, the speedup is more modest (10-30% faster) but still consistent. The only case showing slight regression is when extra_conn_args is an empty dict, where the original's default parameter approach is marginally faster than the or {} fallback.

The optimization is particularly effective when extra_conn_args contains many key-value pairs, as it eliminates the O(n) unpacking operation entirely.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

from __future__ import annotations

# imports
import pytest  # used for our unit tests
from distributed.comm.ws import WSConnector


# Dummy WS class for testing purposes
class WS:
    pass

# Dummy Connector base class for testing purposes
class Connector:
    pass
from distributed.comm.ws import WSConnector

# unit tests

@pytest.fixture
def connector():
    # Fixture to provide a WSConnector instance
    return WSConnector()

# --- Basic Test Cases ---

def test_empty_extra_conn_args_returns_empty_dict(connector):
    # Basic: No extra_conn_args provided, should return empty dict
    codeflash_output = connector._get_connect_args(); result = codeflash_output # 548ns -> 543ns (0.921% faster)

def test_empty_extra_conn_args_dict(connector):
    # Basic: extra_conn_args is an empty dict
    codeflash_output = connector._get_connect_args(extra_conn_args={}); result = codeflash_output # 612ns -> 643ns (4.82% slower)

def test_single_key_value(connector):
    # Basic: extra_conn_args with one key-value pair
    codeflash_output = connector._get_connect_args(extra_conn_args={"foo": "bar"}); result = codeflash_output # 685ns -> 611ns (12.1% faster)

def test_multiple_key_values(connector):
    # Basic: extra_conn_args with multiple key-value pairs
    extra = {"a": 1, "b": 2, "c": 3}
    codeflash_output = connector._get_connect_args(extra_conn_args=extra); result = codeflash_output # 645ns -> 560ns (15.2% faster)

# --- Edge Test Cases ---

def test_extra_conn_args_is_none(connector):
    # Edge: extra_conn_args is None, should treat as empty dict
    codeflash_output = connector._get_connect_args(extra_conn_args=None); result = codeflash_output

def test_extra_conn_args_with_non_string_keys(connector):
    # Edge: keys are not strings
    extra = {1: "one", (2, 3): "tuple"}
    codeflash_output = connector._get_connect_args(extra_conn_args=extra); result = codeflash_output # 900ns -> 563ns (59.9% faster)

def test_extra_conn_args_with_mutable_values(connector):
    # Edge: values are mutable objects
    extra = {"list": [1, 2], "dict": {"x": 10}}
    codeflash_output = connector._get_connect_args(extra_conn_args=extra); result = codeflash_output # 707ns -> 592ns (19.4% faster)

def test_extra_conn_args_with_nested_dicts(connector):
    # Edge: values are nested dicts
    extra = {"nested": {"a": 1, "b": 2}}
    codeflash_output = connector._get_connect_args(extra_conn_args=extra); result = codeflash_output # 632ns -> 550ns (14.9% faster)

def test_extra_conn_args_with_empty_string_key(connector):
    # Edge: empty string as key
    extra = {"": "empty"}
    codeflash_output = connector._get_connect_args(extra_conn_args=extra); result = codeflash_output # 611ns -> 523ns (16.8% faster)

def test_extra_conn_args_with_falsey_keys(connector):
    # Edge: keys are falsey values
    extra = {0: "zero", False: "false"}
    codeflash_output = connector._get_connect_args(extra_conn_args=extra); result = codeflash_output # 603ns -> 503ns (19.9% faster)

def test_extra_conn_args_with_large_integers_and_unusual_types(connector):
    # Edge: very large integer as key, unusual type as value
    extra = {999999999999: object()}
    codeflash_output = connector._get_connect_args(extra_conn_args=extra); result = codeflash_output # 608ns -> 538ns (13.0% faster)

def test_extra_conn_args_with_duplicate_keys(connector):
    # Edge: duplicate keys not possible in dict, but test overwriting
    extra = {"a": 1, "a": 2}
    codeflash_output = connector._get_connect_args(extra_conn_args=extra); result = codeflash_output # 570ns -> 510ns (11.8% faster)

def test_extra_conn_args_with_special_char_keys(connector):
    # Edge: keys with special characters
    extra = {"!@#$%^&*()": "special"}
    codeflash_output = connector._get_connect_args(extra_conn_args=extra); result = codeflash_output # 603ns -> 538ns (12.1% faster)

def test_extra_conn_args_with_boolean_keys(connector):
    # Edge: Boolean keys
    extra = {True: "truthy", False: "falsey"}
    codeflash_output = connector._get_connect_args(extra_conn_args=extra); result = codeflash_output # 636ns -> 491ns (29.5% faster)

def test_extra_conn_args_with_none_key(connector):
    # Edge: None as a key
    extra = {None: "none"}
    codeflash_output = connector._get_connect_args(extra_conn_args=extra); result = codeflash_output # 628ns -> 515ns (21.9% faster)

def test_extra_conn_args_with_tuple_key(connector):
    # Edge: tuple as a key
    extra = {(1, 2): "tuple"}
    codeflash_output = connector._get_connect_args(extra_conn_args=extra); result = codeflash_output # 631ns -> 512ns (23.2% faster)

def test_extra_conn_args_with_bytes_key(connector):
    # Edge: bytes as a key
    extra = {b"bytes": "value"}
    codeflash_output = connector._get_connect_args(extra_conn_args=extra); result = codeflash_output # 591ns -> 458ns (29.0% faster)

def test_extra_conn_args_with_frozenset_key(connector):
    # Edge: frozenset as a key
    extra = {frozenset({1, 2}): "frozenset"}
    codeflash_output = connector._get_connect_args(extra_conn_args=extra); result = codeflash_output # 587ns -> 504ns (16.5% faster)

def test_extra_conn_args_with_large_string_key(connector):
    # Edge: very large string as key
    large_key = "x" * 1000
    extra = {large_key: "large"}
    codeflash_output = connector._get_connect_args(extra_conn_args=extra); result = codeflash_output # 608ns -> 495ns (22.8% faster)

def test_extra_conn_args_with_large_string_value(connector):
    # Edge: very large string as value
    large_value = "y" * 1000
    extra = {"key": large_value}
    codeflash_output = connector._get_connect_args(extra_conn_args=extra); result = codeflash_output # 595ns -> 519ns (14.6% faster)

def test_extra_conn_args_with_unicode_key_and_value(connector):
    # Edge: unicode characters in key and value
    extra = {"ключ": "значение", "😊": "smile"}
    codeflash_output = connector._get_connect_args(extra_conn_args=extra); result = codeflash_output # 651ns -> 491ns (32.6% faster)

def test_extra_conn_args_with_mixed_types(connector):
    # Edge: mixed types in values
    extra = {"int": 1, "float": 2.5, "str": "hello", "bool": True, "none": None}
    codeflash_output = connector._get_connect_args(extra_conn_args=extra); result = codeflash_output # 617ns -> 499ns (23.6% faster)

def test_extra_conn_args_with_additional_kwargs_ignored(connector):
    # Edge: extra kwargs besides extra_conn_args should be ignored
    codeflash_output = connector._get_connect_args(extra_conn_args={"x": 1}, irrelevant=123); result = codeflash_output # 702ns -> 683ns (2.78% faster)

# --- Large Scale Test Cases ---

def test_large_scale_extra_conn_args(connector):
    # Large scale: 1000 key-value pairs
    extra = {f"key_{i}": i for i in range(1000)}
    codeflash_output = connector._get_connect_args(extra_conn_args=extra); result = codeflash_output # 3.19μs -> 506ns (529% faster)

def test_large_scale_extra_conn_args_with_large_values(connector):
    # Large scale: 1000 keys, each value is a large string
    large_value = "x" * 500
    extra = {f"key_{i}": large_value for i in range(1000)}
    codeflash_output = connector._get_connect_args(extra_conn_args=extra); result = codeflash_output # 2.55μs -> 582ns (338% faster)

def test_large_scale_extra_conn_args_with_nested_dicts(connector):
    # Large scale: 1000 keys, each value is a small dict
    extra = {f"key_{i}": {"nested": i} for i in range(1000)}
    codeflash_output = connector._get_connect_args(extra_conn_args=extra); result = codeflash_output # 3.18μs -> 658ns (383% faster)

def test_large_scale_extra_conn_args_with_varied_types(connector):
    # Large scale: 1000 keys, varied types
    extra = {}
    for i in range(1000):
        if i % 5 == 0:
            extra[i] = i
        elif i % 5 == 1:
            extra[i] = str(i)
        elif i % 5 == 2:
            extra[i] = [i]
        elif i % 5 == 3:
            extra[i] = {"x": i}
        else:
            extra[i] = None
    codeflash_output = connector._get_connect_args(extra_conn_args=extra); result = codeflash_output # 2.71μs -> 693ns (290% faster)

def test_large_scale_extra_conn_args_with_large_keys(connector):
    # Large scale: 1000 keys, each key is a large string
    extra = {"k" * (i+1): i for i in range(1000)}
    codeflash_output = connector._get_connect_args(extra_conn_args=extra); result = codeflash_output # 3.61μs -> 600ns (502% faster)

def test_large_scale_extra_conn_args_with_large_mutable_values(connector):
    # Large scale: 1000 keys, each value is a large list
    extra = {f"key_{i}": [i] * 100 for i in range(1000)}
    codeflash_output = connector._get_connect_args(extra_conn_args=extra); result = codeflash_output # 3.65μs -> 636ns (473% faster)

def test_large_scale_extra_conn_args_with_large_tuple_keys(connector):
    # Large scale: 1000 keys, each key is a tuple
    extra = {(i, i+1): i for i in range(1000)}
    codeflash_output = connector._get_connect_args(extra_conn_args=extra); result = codeflash_output # 3.49μs -> 592ns (490% faster)

def test_large_scale_extra_conn_args_with_large_bytes_keys(connector):
    # Large scale: 1000 keys, each key is bytes
    extra = {bytes([i % 256]): i for i in range(1000)}
    codeflash_output = connector._get_connect_args(extra_conn_args=extra); result = codeflash_output # 1.22μs -> 546ns (123% faster)

def test_large_scale_extra_conn_args_with_large_frozenset_keys(connector):
    # Large scale: 1000 keys, each key is a frozenset
    extra = {frozenset({i, i+1}): i for i in range(1000)}
    codeflash_output = connector._get_connect_args(extra_conn_args=extra); result = codeflash_output # 3.50μs -> 578ns (505% faster)

def test_large_scale_extra_conn_args_with_unicode_keys_and_values(connector):
    # Large scale: 1000 keys, unicode strings
    extra = {f"ключ_{i}": f"значение_{i}" for i in range(1000)}
    codeflash_output = connector._get_connect_args(extra_conn_args=extra); result = codeflash_output # 3.52μs -> 586ns (501% faster)
# 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

# imports
import pytest  # used for our unit tests
from distributed.comm.ws import WSConnector

# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

#------------------------------------------------
from distributed.comm.ws import WSConnector

def test_WSConnector__get_connect_args():
    WSConnector._get_connect_args(WSConnector())

🔎 Concolic Coverage Tests and Runtime

Test File::Test Function	Original ⏱️	Optimized ⏱️	Speedup
codeflash_concolic_9h2cxp00/tmpa5dtvs2f/test_concolic_coverage.py::test_WSConnector__get_connect_args	475ns	468ns	1.50%✅

To edit these changes git checkout codeflash/optimize-WSConnector._get_connect_args-mgbsx2br and push.