From 670009ca2273e65ff76a158afe0503d19d46fccc Mon Sep 17 00:00:00 2001
From: Itai Steinherz <itaisteinherz@gmail.com>
Date: Tue, 8 Feb 2022 17:08:14 +0200
Subject: [PATCH] Improve `sock` typing in Paramiko stub (#7156)

---
 stubs/paramiko/paramiko/client.pyi    | 5 +++--
 stubs/paramiko/paramiko/transport.pyi | 7 ++++---
 2 files changed, 7 insertions(+), 5 deletions(-)

diff --git a/stubs/paramiko/paramiko/client.pyi b/stubs/paramiko/paramiko/client.pyi
index 4bbf3c2dc..42fca3991 100644
--- a/stubs/paramiko/paramiko/client.pyi
+++ b/stubs/paramiko/paramiko/client.pyi
@@ -1,4 +1,3 @@
-from socket import socket
 from typing import Iterable, Mapping, NoReturn
 
 from paramiko.channel import Channel, ChannelFile, ChannelStderrFile, ChannelStdinFile
@@ -8,6 +7,8 @@ from paramiko.sftp_client import SFTPClient
 from paramiko.transport import Transport
 from paramiko.util import ClosingContextManager
 
+from .transport import _SocketLike
+
 class SSHClient(ClosingContextManager):
     def __init__(self) -> None: ...
     def load_system_host_keys(self, filename: str | None = ...) -> None: ...
@@ -28,7 +29,7 @@ class SSHClient(ClosingContextManager):
         allow_agent: bool = ...,
         look_for_keys: bool = ...,
         compress: bool = ...,
-        sock: socket | None = ...,
+        sock: _SocketLike | None = ...,
         gss_auth: bool = ...,
         gss_kex: bool = ...,
         gss_deleg_creds: bool = ...,
diff --git a/stubs/paramiko/paramiko/transport.pyi b/stubs/paramiko/paramiko/transport.pyi
index 678a47f93..ea20ff347 100644
--- a/stubs/paramiko/paramiko/transport.pyi
+++ b/stubs/paramiko/paramiko/transport.pyi
@@ -2,7 +2,7 @@ from logging import Logger
 from socket import socket
 from threading import Condition, Event, Lock, Thread
 from types import ModuleType
-from typing import Any, Callable, Iterable, Protocol, Sequence
+from typing import Any, Callable, Iterable, Protocol, Sequence, Union
 
 from paramiko.auth_handler import AuthHandler, _InteractiveCallback
 from paramiko.channel import Channel
@@ -15,6 +15,7 @@ from paramiko.ssh_gss import _SSH_GSSAuth
 from paramiko.util import ClosingContextManager
 
 _Addr = tuple[str, int]
+_SocketLike = Union[str, _Addr, socket, Channel]
 
 class _KexEngine(Protocol):
     def start_kex(self) -> None: ...
@@ -23,7 +24,7 @@ class _KexEngine(Protocol):
 class Transport(Thread, ClosingContextManager):
     active: bool
     hostname: str | None
-    sock: socket
+    sock: socket | Channel
     packetizer: Packetizer
     local_version: str
     remote_version: str
@@ -71,7 +72,7 @@ class Transport(Thread, ClosingContextManager):
     sys: ModuleType
     def __init__(
         self,
-        sock: str | tuple[str, int] | socket,
+        sock: _SocketLike,
         default_window_size: int = ...,
         default_max_packet_size: int = ...,
         gss_kex: bool = ...,