⚡️ Speed up function _init_connection by 20%

[codeflash-ai]

📄 20% (0.20x) speedup for _init_connection in framework/py/flwr/compat/client/app.py

⏱️ Runtime : 4.64 milliseconds → 3.88 milliseconds (best of 156 runs)

📝 Explanation and details

The optimization achieves a 19% speedup by making a simple but impactful change to the parse_address function in common/address.py:

Key Optimization: Efficient String Processing

• Replaced complex string manipulation: Changed raw_host.translate({ord(i): None for i in "[]"}) to raw_host.strip("[]")
• Added IPV6 constant: Defined IPV6: int = 6 at module level to avoid repeated attribute lookups

Why This Improves Performance:
The original .translate() method creates a new dictionary mapping {ord('['): None, ord(']'): None} on every call. Dictionary creation and character-by-character translation is expensive. The optimized .strip("[]") is a native string operation that removes brackets from both ends in a single pass.

Performance Impact Analysis:
From line profiler results, the optimization reduces execution time in the critical parse_address function from 19.19ms to 16.50ms (14% improvement). Since parse_address is called heavily by _init_connection (line 1390 hits), this micro-optimization compounds significantly.

Test Case Performance:
The optimization particularly benefits scenarios with frequent address parsing:

• Large batch tests (test_many_ipv4_addresses, test_many_ipv6_addresses) see the most improvement
• Basic IPv4/IPv6 address parsing becomes more efficient
• Edge cases with invalid formats still handle correctly but faster

The change maintains identical functionality and error handling while replacing an expensive string operation with a more efficient one.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

import pytest
from compat.client.app import _init_connection

# Basic Test Cases

def test_basic_ipv4_default_transport():
    # Test with IPv4 address and default transport
    conn, addr, err_type = _init_connection(None, "127.0.0.1:8080")

def test_basic_ipv6_default_transport():
    # Test with IPv6 address and default transport
    conn, addr, err_type = _init_connection(None, "[::1]:8080")

def test_basic_ipv4_explicit_transport_grpc_bidi():
    # Test with explicit transport type
    conn, addr, err_type = _init_connection("grpc-bidi", "192.168.1.1:5000")

def test_basic_ipv6_explicit_transport_grpc_bidi():
    # Test with explicit transport type and IPv6
    conn, addr, err_type = _init_connection("grpc-bidi", "[2001:db8::1]:1234")

# Edge Test Cases

def test_invalid_port_number_too_high():
    # Port > 65535 should trigger exit (system exit)
    with pytest.raises(SystemExit):
        _init_connection(None, "127.0.0.1:70000")

def test_invalid_port_number_zero():
    # Port == 0 should trigger exit (system exit)
    with pytest.raises(SystemExit):
        _init_connection(None, "127.0.0.1:0")

def test_invalid_port_number_negative():
    # Port < 0 should trigger exit (system exit)
    with pytest.raises(SystemExit):
        _init_connection(None, "127.0.0.1:-1")

def test_missing_port():
    # Missing port should trigger exit (system exit)
    with pytest.raises(SystemExit):
        _init_connection(None, "127.0.0.1")

def test_non_integer_port():
    # Non-integer port should trigger exit (system exit)
    with pytest.raises(SystemExit):
        _init_connection(None, "127.0.0.1:abc")

def test_empty_address():
    # Empty address should trigger exit (system exit)
    with pytest.raises(SystemExit):
        _init_connection(None, "")

def test_invalid_ipv6_format():
    # Invalid IPv6 format should still parse as host, but version is None
    conn, addr, err_type = _init_connection(None, "[invalidv6]:8080")

def test_domain_name():
    # Test with a domain name and valid port
    conn, addr, err_type = _init_connection(None, "flower.ai:8080")

def test_ipv6_without_brackets():
    # IPv6 without brackets should be treated as domain, version=None
    conn, addr, err_type = _init_connection(None, "2001:db8::1:8080")

def test_unknown_transport_type():
    # Unknown transport type should raise ValueError
    with pytest.raises(ValueError):
        _init_connection("unknown-transport", "127.0.0.1:8080")

def test_ipv6_with_extra_brackets():
    # Extra brackets should be stripped by parse_address
    conn, addr, err_type = _init_connection(None, "[[::1]]:8080")

def test_ipv4_with_whitespace():
    # Whitespace should not affect parsing
    conn, addr, err_type = _init_connection(None, " 127.0.0.1 :8080 ")

# Large Scale Test Cases

def test_many_ipv4_addresses():
    # Test with many IPv4 addresses to ensure scalability
    for i in range(1, 100):
        addr_str = f"10.0.0.{i}:5000"
        conn, addr, err_type = _init_connection(None, addr_str)

def test_many_ipv6_addresses():
    # Test with many IPv6 addresses to ensure scalability
    for i in range(1, 100):
        addr_str = f"[2001:db8::{i}]:6000"
        conn, addr, err_type = _init_connection(None, addr_str)

def test_large_port_numbers_within_range():
    # Test with large, but valid port numbers
    for port in range(65000, 65536):
        addr_str = f"127.0.0.1:{port}"
        conn, addr, err_type = _init_connection(None, addr_str)

def test_long_domain_name():
    # Test with a long domain name
    domain = "a" * 200 + ".flower.ai"
    addr_str = f"{domain}:8080"
    conn, addr, err_type = _init_connection(None, addr_str)

def test_large_number_of_calls():
    # Call _init_connection multiple times to ensure no resource leaks
    results = []
    for _ in range(500):
        conn, addr, err_type = _init_connection(None, "127.0.0.1:8080")
        results.append((conn, addr, err_type))
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
import builtins
from contextlib import AbstractContextManager
from typing import Callable, Optional, Union

# imports
import pytest
from compat.client.app import _init_connection
from cryptography.hazmat.primitives.asymmetric import ec

# Basic Test Cases

def test_ipv4_address_default_transport():
    # Test with IPv4 address, default transport (None)
    conn, addr, err_type = _init_connection(None, "127.0.0.1:8080")

def test_ipv4_address_explicit_grpc_bidi():
    # Test with IPv4 address, explicit grpc-bidi transport
    conn, addr, err_type = _init_connection("grpc-bidi", "192.168.1.10:5000")

def test_ipv6_address_default_transport():
    # Test with IPv6 address, default transport
    conn, addr, err_type = _init_connection(None, "[::1]:8080")

def test_ipv6_address_explicit_grpc_bidi():
    # Test with IPv6 address, explicit grpc-bidi
    conn, addr, err_type = _init_connection("grpc-bidi", "[2001:db8::1]:12345")

def test_domain_address():
    # Test with domain name address
    conn, addr, err_type = _init_connection(None, "example.com:443")

# Edge Test Cases

def test_invalid_port_too_low():
    # Port number too low (0)
    with pytest.raises(SystemExit) as excinfo:
        _init_connection(None, "127.0.0.1:0")

def test_invalid_port_too_high():
    # Port number too high (65536)
    with pytest.raises(SystemExit) as excinfo:
        _init_connection(None, "127.0.0.1:65536")

def test_missing_port():
    # Missing port number
    with pytest.raises(SystemExit) as excinfo:
        _init_connection(None, "127.0.0.1")

def test_invalid_address_format():
    # Completely invalid address
    with pytest.raises(SystemExit) as excinfo:
        _init_connection(None, "not-an-address")

def test_empty_address():
    # Empty string as address
    with pytest.raises(SystemExit) as excinfo:
        _init_connection(None, "")

def test_ipv6_without_brackets():
    # IPv6 address without brackets (should be parsed as domain)
    conn, addr, err_type = _init_connection(None, "2001:db8::1:8080")

def test_unknown_transport_type():
    # Unknown transport type should raise ValueError
    with pytest.raises(ValueError):
        _init_connection("unknown-transport", "127.0.0.1:8080")

def test_ipv6_with_port_and_domain():
    # Domain that looks like IPv6 with port (should be parsed as domain)
    conn, addr, err_type = _init_connection(None, "[my-domain.com]:9000")

def test_transport_case_sensitivity():
    # Should be case-sensitive: "GRPC-BIDI" is not valid
    with pytest.raises(ValueError):
        _init_connection("GRPC-BIDI", "127.0.0.1:8080")

def test_port_as_string():
    # Port as string (should be parsed correctly)
    conn, addr, err_type = _init_connection(None, "localhost:12345")

# Large Scale Test Cases

@pytest.mark.parametrize("i", range(1, 100, 13))
def test_many_ipv4_addresses(i):
    # Test a range of valid IPv4 addresses and ports
    ip = f"10.0.0.{i}"
    port = 10000 + i
    addr_str = f"{ip}:{port}"
    conn, addr, err_type = _init_connection(None, addr_str)

@pytest.mark.parametrize("i", range(1, 100, 13))
def test_many_ipv6_addresses(i):
    # Test a range of valid IPv6 addresses and ports
    ip = f"2001:db8::{i}"
    port = 20000 + i
    addr_str = f"[{ip}]:{port}"
    conn, addr, err_type = _init_connection(None, addr_str)

def test_large_batch_of_addresses():
    # Test a batch of 100 addresses to check scalability
    results = []
    for i in range(1, 101):
        ip = f"172.16.0.{i}"
        port = 30000 + i
        addr_str = f"{ip}:{port}"
        conn, addr, err_type = _init_connection(None, addr_str)
        results.append((addr, conn, err_type))

def test_large_batch_of_invalid_ports():
    # Test a batch of invalid ports for error handling
    for port in [0, 65536, 99999]:
        with pytest.raises(SystemExit):
            _init_connection(None, f"10.10.10.10:{port}")

def test_large_batch_of_invalid_addresses():
    # Test a batch of invalid addresses
    invalids = ["", ":", "::::", "noport", "[]:1234", "[::1]", "127.0.0.1:", "localhost:"]
    for addr in invalids:
        with pytest.raises(SystemExit):
            _init_connection(None, addr)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

To edit these changes git checkout codeflash/optimize-_init_connection-mhc9akh5 and push.