feat: add callback for users to customize socket creation (#147)

Co-authored-by: Kieren <kjander0>

Author: kjander0

src/aiohappyeyeballs/__init__.py

@@ -1,11 +1,12 @@
 __version__ = "2.4.8"
 
 from .impl import start_connection
-from .types import AddrInfoType
+from .types import AddrInfoType, SocketFactoryType
 from .utils import addr_to_addr_infos, pop_addr_infos_interleave, remove_addr_infos
 
 __all__ = (
     "AddrInfoType",
+    "SocketFactoryType",
     "addr_to_addr_infos",
     "pop_addr_infos_interleave",
     "remove_addr_infos",

src/aiohappyeyeballs/impl.py

@@ -9,7 +9,7 @@
 from typing import List, Optional, Sequence, Set, Union
 
 from . import _staggered
-from .types import AddrInfoType
+from .types import AddrInfoType, SocketFactoryType
 
 
 async def start_connection(
@@ -19,6 +19,7 @@ async def start_connection(
     happy_eyeballs_delay: Optional[float] = None,
     interleave: Optional[int] = None,
     loop: Optional[asyncio.AbstractEventLoop] = None,
+    socket_factory: Optional[SocketFactoryType] = None,
 ) -> socket.socket:
     """
     Connect to a TCP server.
@@ -70,7 +71,12 @@ async def start_connection(
         for addrinfo in addr_infos:
             try:
                 sock = await _connect_sock(
-                    current_loop, exceptions, addrinfo, local_addr_infos
+                    current_loop,
+                    exceptions,
+                    addrinfo,
+                    local_addr_infos,
+                    None,
+                    socket_factory,
                 )
                 break
             except (RuntimeError, OSError):
@@ -87,6 +93,7 @@ async def start_connection(
                         addrinfo,
                         local_addr_infos,
                         open_sockets,
+                        socket_factory,
                     )
                     for addrinfo in addr_infos
                 ),
@@ -153,6 +160,7 @@ async def _connect_sock(
     addr_info: AddrInfoType,
     local_addr_infos: Optional[Sequence[AddrInfoType]] = None,
     open_sockets: Optional[Set[socket.socket]] = None,
+    socket_factory: Optional[SocketFactoryType] = None,
 ) -> socket.socket:
     """
     Create, bind and connect one socket.
@@ -169,7 +177,10 @@ async def _connect_sock(
     family, type_, proto, _, address = addr_info
     sock = None
     try:
-        sock = socket.socket(family=family, type=type_, proto=proto)
+        if socket_factory is not None:
+            sock = socket_factory(addr_info)
+        else:
+            sock = socket.socket(family=family, type=type_, proto=proto)
         if open_sockets is not None:
             open_sockets.add(sock)
         sock.setblocking(False)

src/aiohappyeyeballs/types.py

@@ -1,6 +1,7 @@
 """Types for aiohappyeyeballs."""
 
 import socket
+from collections.abc import Callable
 from typing import Tuple, Union
 
 AddrInfoType = Tuple[
@@ -10,3 +11,5 @@
     str,
     Tuple,  # type: ignore[type-arg]
 ]
+
+SocketFactoryType = Callable[[AddrInfoType], socket.socket]

tests/test_impl.py

@@ -6,7 +6,13 @@
 
 import pytest
 
-from aiohappyeyeballs import AddrInfoType, _staggered, impl, start_connection
+from aiohappyeyeballs import (
+    AddrInfoType,
+    SocketFactoryType,
+    _staggered,
+    impl,
+    start_connection,
+)
 
 
 def mock_socket_module():
@@ -136,6 +142,33 @@ def _socket(*args, **kw):
         )
 
 
+@pytest.mark.asyncio
+@patch_socket
+async def test_single_addr_socket_factory(m_socket: ModuleType) -> None:
+    mock_socket = mock.MagicMock(
+        family=socket.AF_INET,
+        type=socket.SOCK_STREAM,
+        proto=socket.IPPROTO_TCP,
+        fileno=mock.MagicMock(return_value=1),
+    )
+
+    def factory(addr_info: AddrInfoType) -> socket.socket:
+        return mock_socket
+
+    addr_info = [
+        (
+            socket.AF_INET,
+            socket.SOCK_STREAM,
+            socket.IPPROTO_TCP,
+            "",
+            ("107.6.106.82", 80),
+        )
+    ]
+    loop = asyncio.get_running_loop()
+    with mock.patch.object(loop, "sock_connect", return_value=None):
+        assert await start_connection(addr_info, socket_factory=factory) == mock_socket
+
+
 @pytest.mark.asyncio
 @patch_socket
 async def test_multiple_addr_success_second_one(
@@ -201,6 +234,7 @@ async def _connect_sock(
         addr_info: AddrInfoType,
         local_addr_infos: Optional[Sequence[AddrInfoType]] = None,
         sockets: Optional[Set[socket.socket]] = None,
+        socket_factory: Optional[SocketFactoryType] = None,
     ) -> socket.socket:
         await finish
         sock = _socket()
@@ -293,6 +327,54 @@ def _socket(*args, **kw):
         )
 
 
+@pytest.mark.asyncio
+@patch_socket
+async def test_happy_eyeballs_socket_factory(
+    m_socket: ModuleType,
+) -> None:
+    mock_socket = mock.MagicMock(
+        family=socket.AF_INET,
+        type=socket.SOCK_STREAM,
+        proto=socket.IPPROTO_TCP,
+        fileno=mock.MagicMock(return_value=1),
+    )
+
+    idx = -1
+    errors = ["err1", "err2"]
+
+    def factory(addr_info: AddrInfoType) -> socket.socket:
+        nonlocal idx, errors
+        idx += 1
+        if idx == 1:
+            raise OSError(5, errors[idx])
+        return mock_socket
+
+    addr_info = [
+        (
+            socket.AF_INET,
+            socket.SOCK_STREAM,
+            socket.IPPROTO_TCP,
+            "",
+            ("107.6.106.82", 80),
+        ),
+        (
+            socket.AF_INET,
+            socket.SOCK_STREAM,
+            socket.IPPROTO_TCP,
+            "",
+            ("107.6.106.83", 80),
+        ),
+    ]
+    loop = asyncio.get_running_loop()
+    with mock.patch.object(loop, "sock_connect", return_value=None):
+        assert (
+            await start_connection(
+                addr_info, happy_eyeballs_delay=0.3, socket_factory=factory
+            )
+            == mock_socket
+        )
+
+
 @pytest.mark.asyncio
 @patch_socket
 async def test_multiple_addr_all_fail_happy_eyeballs(