Add urllib3.contrib.socks; improve urllib3.connectionpool (#8457)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Alex Waygood <[email protected]>

Author: 3 people

stubs/urllib3/@tests/requirements-stubtest.txt

@@ -0,0 +1 @@
+PySocks>=1.5.6,<2.0,!=1.5.7

stubs/urllib3/@tests/stubtest_allowlist.txt

@@ -1,3 +1,5 @@
+# TODO: remove ResponseCls ignore when https://github.com/python/mypy/issues/13316 is closed
+urllib3.HTTPConnectionPool.ResponseCls
 urllib3.HTTPConnectionPool.__init__
 urllib3.HTTPConnectionPool.urlopen
 urllib3.HTTPSConnectionPool.__init__
@@ -17,6 +19,8 @@ urllib3.connection.VerifiedHTTPSConnection.__init__
 urllib3.connection.VerifiedHTTPSConnection.set_cert
 urllib3.connectionpool.ConnectionError
 urllib3.connectionpool.HTTPConnection.request
+# TODO: remove ResponseCls ignore when https://github.com/python/mypy/issues/13316 is closed
+urllib3.connectionpool.HTTPConnectionPool.ResponseCls
 urllib3.connectionpool.HTTPConnectionPool.__init__
 urllib3.connectionpool.HTTPConnectionPool.urlopen
 urllib3.connectionpool.HTTPSConnection.__init__

stubs/urllib3/urllib3/connectionpool.pyi

@@ -1,9 +1,15 @@
-from typing import Any
+import queue
+from _typeshed import Self
+from collections.abc import Mapping
+from logging import Logger
+from types import TracebackType
+from typing import Any, ClassVar
+from typing_extensions import Literal, TypeAlias
 
 from . import connection, exceptions, request, response
 from .connection import BaseSSLError as BaseSSLError, ConnectionError as ConnectionError, HTTPException as HTTPException
 from .packages import ssl_match_hostname
-from .util import connection as _connection, retry, timeout, url
+from .util import Url, connection as _connection, queue as urllib3queue, retry, timeout, url
 
 ClosedPoolError = exceptions.ClosedPoolError
 ProtocolError = exceptions.ProtocolError
@@ -29,48 +35,54 @@ Retry = retry.Retry
 Timeout = timeout.Timeout
 get_host = url.get_host
 
+_Timeout: TypeAlias = Timeout | float
+_Retries: TypeAlias = Retry | bool | int
+
 xrange: Any
-log: Any
+log: Logger
 
 class ConnectionPool:
-    scheme: Any
-    QueueCls: Any
-    host: Any
-    port: Any
-    def __init__(self, host, port=...) -> None: ...
-    def __enter__(self): ...
-    def __exit__(self, exc_type, exc_val, exc_tb): ...
-    def close(self): ...
+    scheme: ClassVar[str | None]
+    QueueCls: ClassVar[type[queue.Queue[Any]]]
+    host: str
+    port: int | None
+    def __init__(self, host: str, port: int | None = ...) -> None: ...
+    def __enter__(self: Self) -> Self: ...
+    def __exit__(
+        self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: TracebackType | None
+    ) -> Literal[False]: ...
+    def close(self) -> None: ...
 
 class HTTPConnectionPool(ConnectionPool, RequestMethods):
-    scheme: Any
-    ConnectionCls: Any
-    strict: Any
-    timeout: Any
-    retries: Any
-    pool: Any
-    block: Any
-    proxy: Any
-    proxy_headers: Any
-    num_connections: Any
-    num_requests: Any
+    scheme: ClassVar[str]
+    ConnectionCls: ClassVar[type[HTTPConnection | HTTPSConnection]]
+    ResponseCls: ClassVar[type[HTTPResponse]]
+    strict: bool
+    timeout: _Timeout
+    retries: _Retries | None
+    pool: urllib3queue.LifoQueue | None
+    block: bool
+    proxy: Url | None
+    proxy_headers: Mapping[str, str]
+    num_connections: int
+    num_requests: int
     conn_kw: Any
     def __init__(
         self,
-        host,
-        port=...,
-        strict=...,
-        timeout=...,
-        maxsize=...,
-        block=...,
-        headers=...,
-        retries=...,
-        _proxy=...,
-        _proxy_headers=...,
+        host: str,
+        port: int | None = ...,
+        strict: bool = ...,
+        timeout: _Timeout = ...,
+        maxsize: int = ...,
+        block: bool = ...,
+        headers: Mapping[str, str] | None = ...,
+        retries: _Retries | None = ...,
+        _proxy: Url | None = ...,
+        _proxy_headers: Mapping[str, str] | None = ...,
         **conn_kw,
     ) -> None: ...
-    def close(self): ...
-    def is_same_host(self, url): ...
+    def close(self) -> None: ...
+    def is_same_host(self, url: str) -> bool: ...
     def urlopen(
         self,
         method,
@@ -87,35 +99,33 @@ class HTTPConnectionPool(ConnectionPool, RequestMethods):
     ): ...
 
 class HTTPSConnectionPool(HTTPConnectionPool):
-    scheme: Any
-    ConnectionCls: Any
-    key_file: Any
-    cert_file: Any
-    cert_reqs: Any
-    ca_certs: Any
-    ssl_version: Any
-    assert_hostname: Any
-    assert_fingerprint: Any
+    key_file: str | None
+    cert_file: str | None
+    cert_reqs: int | str | None
+    ca_certs: str | None
+    ssl_version: int | str | None
+    assert_hostname: str | Literal[False] | None
+    assert_fingerprint: str | None
     def __init__(
         self,
-        host,
-        port=...,
-        strict=...,
-        timeout=...,
-        maxsize=...,
-        block=...,
-        headers=...,
-        retries=...,
-        _proxy=...,
-        _proxy_headers=...,
-        key_file=...,
-        cert_file=...,
-        cert_reqs=...,
-        ca_certs=...,
-        ssl_version=...,
-        assert_hostname=...,
-        assert_fingerprint=...,
+        host: str,
+        port: int | None = ...,
+        strict: bool = ...,
+        timeout: _Timeout = ...,
+        maxsize: int = ...,
+        block: bool = ...,
+        headers: Mapping[str, str] | None = ...,
+        retries: _Retries | None = ...,
+        _proxy: Url | None = ...,
+        _proxy_headers: Mapping[str, str] | None = ...,
+        key_file: str | None = ...,
+        cert_file: str | None = ...,
+        cert_reqs: int | str | None = ...,
+        ca_certs: str | None = ...,
+        ssl_version: int | str | None = ...,
+        assert_hostname: str | Literal[False] | None = ...,
+        assert_fingerprint: str | None = ...,
         **conn_kw,
     ) -> None: ...
 
-def connection_from_url(url, **kw): ...
+def connection_from_url(url: str, **kw) -> HTTPConnectionPool: ...

stubs/urllib3/urllib3/contrib/socks.pyi

@@ -0,0 +1,45 @@
+from collections.abc import Mapping
+from typing import ClassVar
+from typing_extensions import TypedDict
+
+from ..connection import HTTPConnection, HTTPSConnection
+from ..connectionpool import HTTPConnectionPool, HTTPSConnectionPool
+from ..poolmanager import PoolManager
+
+class _TYPE_SOCKS_OPTIONS(TypedDict):
+    socks_version: int
+    proxy_host: str | None
+    proxy_port: str | None
+    username: str | None
+    password: str | None
+    rdns: bool
+
+class SOCKSConnection(HTTPConnection):
+    def __init__(self, _socks_options: _TYPE_SOCKS_OPTIONS, *args, **kwargs) -> None: ...
+
+class SOCKSHTTPSConnection(SOCKSConnection, HTTPSConnection): ...
+
+class SOCKSHTTPConnectionPool(HTTPConnectionPool):
+    ConnectionCls: ClassVar[type[SOCKSConnection]]
+
+class SOCKSHTTPSConnectionPool(HTTPSConnectionPool):
+    ConnectionCls: ClassVar[type[SOCKSHTTPSConnection]]
+
+class _ConnectionPoolClasses(TypedDict):
+    http: type[SOCKSHTTPConnectionPool]
+    https: type[SOCKSHTTPSConnectionPool]
+
+class SOCKSProxyManager(PoolManager):
+    # has a class-level default, but is overridden on instances, so not a ClassVar
+    pool_classes_by_scheme: _ConnectionPoolClasses
+    proxy_url: str
+
+    def __init__(
+        self,
+        proxy_url: str,
+        username: str | None = ...,
+        password: str | None = ...,
+        num_pools: int = ...,
+        headers: Mapping[str, str] | None = ...,
+        **connection_pool_kw,
+    ) -> None: ...

stubs/urllib3/urllib3/util/queue.pyi

@@ -0,0 +1,4 @@
+from queue import Queue
+from typing import Any
+
+class LifoQueue(Queue[Any]): ...