⚡️ Speed up method WSSConnector._get_connect_args by 76%

[codeflash-ai]

📄 76% (0.76x) speedup for WSSConnector._get_connect_args in distributed/comm/ws.py

⏱️ Runtime : 109 microseconds → 62.1 microseconds (best of 148 runs)

📝 Explanation and details

The optimized code achieves a 76% speedup by eliminating expensive dictionary unpacking operations and reducing redundant lookups. Here are the key changes:

1. Avoided Dictionary Unpacking for Large Dictionaries

• Original: wss_args = {"ssl_options": ..., **connection_args.get("extra_conn_args", {})}
• Optimized: Uses dict(extra_conn_args) followed by direct assignment wss_args["ssl_options"] = ssl_options
• Why faster: Dictionary unpacking (**dict) creates temporary objects and iterates through all key-value pairs, while dict() constructor is more efficient for copying

2. Eliminated Redundant Dictionary Lookups

• Original: Multiple calls to connection_args.get() within dictionary literals
• Optimized: Pre-fetches values (ssl_options, extra_conn_args, server_hostname) once
• Why faster: Reduces function call overhead and hash table lookups

3. Optimized Headers Merging

• Original: {**wss_args.get("headers", {}), **{"Host": connection_args["server_hostname"]}}
• Optimized: Direct copying with dict(headers) and assignment headers["Host"] = server_hostname
• Why faster: Avoids creating multiple temporary dictionaries during unpacking

Performance Benefits by Test Case:

• Large dictionaries (1000+ keys): 223-259% faster - the unpacking avoidance provides massive gains
• Basic cases: 4-17% faster - reduced lookup overhead helps even small operations
• Header merging: 6-35% faster - direct assignment beats unpacking for header updates

The optimizations are particularly effective for large extra_conn_args dictionaries, which is common in distributed computing scenarios where connection configurations can be extensive.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

from __future__ import annotations

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


# Dummy parent classes for context (since WSConnector and WSS are not defined)
class WSConnector:
    pass

class WSS:
    pass
from distributed.comm.ws import WSSConnector

# unit tests

# Helper fixture to instantiate WSSConnector
@pytest.fixture
def connector():
    return WSSConnector()

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

def test_empty_args(connector):
    """Test with no arguments: should only include ssl_options=None."""
    codeflash_output = connector._get_connect_args(); result = codeflash_output # 777ns -> 725ns (7.17% faster)

def test_only_ssl_context(connector):
    """Test with only ssl_context provided."""
    ctx = object()
    codeflash_output = connector._get_connect_args(ssl_context=ctx); result = codeflash_output # 964ns -> 921ns (4.67% faster)

def test_only_extra_conn_args(connector):
    """Test with only extra_conn_args provided."""
    extra = {"foo": "bar", "baz": 42}
    codeflash_output = connector._get_connect_args(extra_conn_args=extra); result = codeflash_output # 933ns -> 1.03μs (9.24% slower)

def test_ssl_context_and_extra_conn_args(connector):
    """Test with both ssl_context and extra_conn_args."""
    ctx = object()
    extra = {"foo": "bar", "baz": 42}
    codeflash_output = connector._get_connect_args(ssl_context=ctx, extra_conn_args=extra); result = codeflash_output # 1.02μs -> 1.05μs (2.87% slower)

def test_server_hostname_only(connector):
    """Test with only server_hostname provided."""
    host = "myserver.com"
    codeflash_output = connector._get_connect_args(server_hostname=host); result = codeflash_output # 1.32μs -> 1.12μs (17.9% faster)

def test_server_hostname_and_extra_headers(connector):
    """Test with server_hostname and extra_conn_args containing headers."""
    host = "example.org"
    extra = {"headers": {"X-Foo": "bar", "X-Bar": "baz"}}
    codeflash_output = connector._get_connect_args(server_hostname=host, extra_conn_args=extra); result = codeflash_output # 1.44μs -> 1.32μs (8.86% faster)

def test_all_arguments(connector):
    """Test with all arguments provided."""
    ctx = "ssl_ctx"
    host = "host.test"
    extra = {"foo": "bar", "headers": {"X-Test": "yes"}}
    codeflash_output = connector._get_connect_args(ssl_context=ctx, server_hostname=host, extra_conn_args=extra); result = codeflash_output # 1.50μs -> 1.36μs (10.5% faster)

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

def test_extra_conn_args_none(connector):
    """Test with extra_conn_args=None (should not fail)."""
    codeflash_output = connector._get_connect_args(extra_conn_args=None); result = codeflash_output

def test_extra_conn_args_empty_dict(connector):
    """Test with extra_conn_args={} (empty dict)."""
    codeflash_output = connector._get_connect_args(extra_conn_args={}); result = codeflash_output # 1.17μs -> 892ns (31.1% faster)


def test_headers_missing_in_extra(connector):
    """Test with extra_conn_args not containing headers, but server_hostname present."""
    host = "noheaders.com"
    extra = {"foo": "bar"}
    codeflash_output = connector._get_connect_args(server_hostname=host, extra_conn_args=extra); result = codeflash_output # 1.43μs -> 1.61μs (11.3% slower)

def test_headers_overlap(connector):
    """Test with extra_conn_args containing headers with Host key, server_hostname should override."""
    host = "override.com"
    extra = {"headers": {"Host": "oldhost.com", "X-Header": "val"}}
    codeflash_output = connector._get_connect_args(server_hostname=host, extra_conn_args=extra); result = codeflash_output # 1.37μs -> 1.38μs (0.073% slower)

def test_none_server_hostname(connector):
    """Test with server_hostname=None (should not add headers)."""
    extra = {"foo": "bar"}
    codeflash_output = connector._get_connect_args(server_hostname=None, extra_conn_args=extra); result = codeflash_output # 953ns -> 987ns (3.44% slower)

def test_empty_string_server_hostname(connector):
    """Test with server_hostname='' (should not add headers)."""
    extra = {"foo": "bar"}
    codeflash_output = connector._get_connect_args(server_hostname="", extra_conn_args=extra); result = codeflash_output # 987ns -> 1.02μs (3.52% slower)


def test_extra_conn_args_with_non_string_keys(connector):
    """Test with extra_conn_args containing non-string keys."""
    extra = {42: "answer", None: "noneval"}
    codeflash_output = connector._get_connect_args(extra_conn_args=extra); result = codeflash_output # 989ns -> 1.38μs (28.1% slower)

def test_extra_conn_args_overwrites_ssl_options(connector):
    """Test if extra_conn_args contains 'ssl_options', it overwrites the key."""
    ctx = "ssl_ctx"
    extra = {"ssl_options": "other_ctx"}
    codeflash_output = connector._get_connect_args(ssl_context=ctx, extra_conn_args=extra); result = codeflash_output # 958ns -> 1.09μs (12.2% slower)

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

def test_large_extra_conn_args(connector):
    """Test with a large number of extra_conn_args."""
    size = 1000
    extra = {f"key{i}": f"value{i}" for i in range(size)}
    codeflash_output = connector._get_connect_args(extra_conn_args=extra); result = codeflash_output # 12.4μs -> 3.82μs (223% faster)
    for i in range(size):
        pass

def test_large_headers_merge(connector):
    """Test with large headers in extra_conn_args and server_hostname."""
    size = 500
    extra_headers = {f"Header{i}": f"Val{i}" for i in range(size)}
    extra = {"headers": extra_headers}
    host = "bighost.com"
    codeflash_output = connector._get_connect_args(server_hostname=host, extra_conn_args=extra); result = codeflash_output # 2.97μs -> 2.71μs (9.69% faster)
    for i in range(size):
        pass

def test_large_all_arguments(connector):
    """Test with large extra_conn_args and server_hostname and ssl_context."""
    size = 900
    extra = {f"k{i}": f"v{i}" for i in range(size)}
    extra["headers"] = {f"H{i}": f"V{i}" for i in range(100)}
    host = "scalebig.com"
    ctx = "bigssl"
    codeflash_output = connector._get_connect_args(ssl_context=ctx, server_hostname=host, extra_conn_args=extra); result = codeflash_output # 12.4μs -> 4.52μs (174% faster)
    for i in range(size):
        if f"k{i}" != "headers":
            pass
    for i in range(100):
        pass
# 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 WSSConnector


class WSConnector:
    pass

class WSS:
    pass
from distributed.comm.ws import WSSConnector

# unit tests

# ---- Basic Test Cases ----

def test_basic_ssl_context_only():
    # Only ssl_context provided
    connector = WSSConnector()
    codeflash_output = connector._get_connect_args(ssl_context="sslctx"); result = codeflash_output # 1.02μs -> 909ns (12.2% faster)

def test_basic_extra_conn_args_only():
    # Only extra_conn_args provided
    connector = WSSConnector()
    extra = {"foo": "bar", "baz": 123}
    codeflash_output = connector._get_connect_args(extra_conn_args=extra); result = codeflash_output # 900ns -> 999ns (9.91% slower)

def test_basic_ssl_and_extra():
    # Both ssl_context and extra_conn_args provided
    connector = WSSConnector()
    extra = {"foo": "bar"}
    codeflash_output = connector._get_connect_args(ssl_context="sslctx", extra_conn_args=extra); result = codeflash_output # 981ns -> 1.01μs (3.06% slower)

def test_basic_server_hostname_only():
    # Only server_hostname provided
    connector = WSSConnector()
    codeflash_output = connector._get_connect_args(server_hostname="myhost.com"); result = codeflash_output # 1.33μs -> 1.13μs (17.5% faster)

def test_basic_all_args():
    # All arguments provided
    connector = WSSConnector()
    extra = {"foo": "bar", "headers": {"X-Token": "abc"}}
    codeflash_output = connector._get_connect_args(
        ssl_context="sslctx",
        extra_conn_args=extra,
        server_hostname="myhost.com"
    ); result = codeflash_output # 1.53μs -> 1.32μs (15.4% faster)

# ---- Edge Test Cases ----

def test_edge_empty_extra_conn_args():
    # extra_conn_args is empty dict
    connector = WSSConnector()
    codeflash_output = connector._get_connect_args(extra_conn_args={}); result = codeflash_output # 859ns -> 945ns (9.10% slower)

def test_edge_none_extra_conn_args():
    # extra_conn_args is None
    connector = WSSConnector()
    codeflash_output = connector._get_connect_args(extra_conn_args=None); result = codeflash_output

def test_edge_headers_in_extra_conn_args_and_server_hostname():
    # extra_conn_args has headers, server_hostname provided
    connector = WSSConnector()
    extra = {"headers": {"X-Token": "abc", "Host": "wronghost.com"}}
    codeflash_output = connector._get_connect_args(
        extra_conn_args=extra,
        server_hostname="correcthost.com"
    ); result = codeflash_output # 1.71μs -> 1.27μs (34.6% faster)

def test_edge_headers_not_dict_in_extra_conn_args():
    # headers in extra_conn_args is not a dict
    connector = WSSConnector()
    extra = {"headers": "notadict"}
    codeflash_output = connector._get_connect_args(
        extra_conn_args=extra,
        server_hostname="myhost.com"
    ); result = codeflash_output

def test_edge_server_hostname_empty_string():
    # server_hostname is empty string
    connector = WSSConnector()
    extra = {"headers": {"X-Token": "abc"}}
    codeflash_output = connector._get_connect_args(
        extra_conn_args=extra,
        server_hostname=""
    ); result = codeflash_output # 1.28μs -> 1.33μs (3.97% slower)

def test_edge_server_hostname_none():
    # server_hostname is None
    connector = WSSConnector()
    extra = {"headers": {"X-Token": "abc"}}
    codeflash_output = connector._get_connect_args(
        extra_conn_args=extra,
        server_hostname=None
    ); result = codeflash_output # 1.02μs -> 1.02μs (0.390% slower)

def test_edge_extra_conn_args_overwrites_ssl_options():
    # extra_conn_args contains ssl_options, should overwrite
    connector = WSSConnector()
    extra = {"ssl_options": "override"}
    codeflash_output = connector._get_connect_args(
        ssl_context="original",
        extra_conn_args=extra
    ); result = codeflash_output # 972ns -> 1.06μs (8.56% slower)





def test_edge_extra_conn_args_has_non_string_keys():
    # extra_conn_args has non-string keys
    connector = WSSConnector()
    extra = {1: "one", None: "none"}
    codeflash_output = connector._get_connect_args(extra_conn_args=extra); result = codeflash_output # 1.14μs -> 1.44μs (20.8% slower)

def test_edge_headers_is_none_in_extra_conn_args():
    # headers in extra_conn_args is None
    connector = WSSConnector()
    extra = {"headers": None}
    codeflash_output = connector._get_connect_args(
        extra_conn_args=extra,
        server_hostname="host.com"
    ); result = codeflash_output

# ---- Large Scale Test Cases ----

def test_large_extra_conn_args_many_keys():
    # Large number of keys in extra_conn_args
    connector = WSSConnector()
    extra = {f"key{i}": i for i in range(1000)}
    codeflash_output = connector._get_connect_args(
        ssl_context="sslctx",
        extra_conn_args=extra
    ); result = codeflash_output # 12.6μs -> 3.51μs (259% faster)
    for i in range(1000):
        pass

def test_large_headers_many_keys():
    # Large number of keys in headers
    connector = WSSConnector()
    headers = {f"X-Key-{i}": f"val{i}" for i in range(1000)}
    extra = {"headers": headers}
    codeflash_output = connector._get_connect_args(
        extra_conn_args=extra,
        server_hostname="bighost.com"
    ); result = codeflash_output # 4.39μs -> 4.10μs (6.90% faster)
    # Host should be added to headers
    for i in range(1000):
        pass

def test_large_headers_merge_and_override():
    # Large headers, Host already present, should be overridden
    connector = WSSConnector()
    headers = {f"X-Key-{i}": f"val{i}" for i in range(999)}
    headers["Host"] = "oldhost.com"
    extra = {"headers": headers}
    codeflash_output = connector._get_connect_args(
        extra_conn_args=extra,
        server_hostname="newhost.com"
    ); result = codeflash_output # 4.27μs -> 4.03μs (5.99% faster)
    for i in range(999):
        pass

def test_large_extra_conn_args_and_headers():
    # Large extra_conn_args and headers together
    connector = WSSConnector()
    extra = {f"key{i}": i for i in range(500)}
    headers = {f"X-Key-{i}": f"val{i}" for i in range(500)}
    extra["headers"] = headers
    codeflash_output = connector._get_connect_args(
        ssl_context="sslctx",
        extra_conn_args=extra,
        server_hostname="host.com"
    ); result = codeflash_output # 9.66μs -> 4.00μs (142% faster)
    for i in range(500):
        pass

def test_large_extra_conn_args_with_non_string_keys():
    # Large extra_conn_args with non-string keys
    connector = WSSConnector()
    extra = {i: f"val{i}" for i in range(1000)}
    codeflash_output = connector._get_connect_args(extra_conn_args=extra); result = codeflash_output # 9.63μs -> 3.34μs (188% faster)
    for i in range(1000):
        pass
# 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 WSSConnector

def test_WSSConnector__get_connect_args():
    WSSConnector._get_connect_args(WSSConnector())

🔎 Concolic Coverage Tests and Runtime

Test File::Test Function	Original ⏱️	Optimized ⏱️	Speedup
codeflash_concolic_9h2cxp00/tmp9f5x08do/test_concolic_coverage.py::test_WSSConnector__get_connect_args	754ns	660ns	14.2%✅

To edit these changes git checkout codeflash/optimize-WSSConnector._get_connect_args-mgbt3lz7 and push.