From ea465204733fdc77a5d8a72dc2462cab9242f385 Mon Sep 17 00:00:00 2001 From: Pieter-Jan Briers Date: Tue, 17 Oct 2023 14:04:47 +0200 Subject: [PATCH 1/2] Use Self TypeVar for SSHConnection __aenter__ This makes it so that when using a subclass as a context manager like so: async with asyncssh.connect(...) as con: the con variable has the type SSHClientConnection instead of simply SSHConnection. --- asyncssh/connection.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/asyncssh/connection.py b/asyncssh/connection.py index 3aaa70c..9038ac4 100644 --- a/asyncssh/connection.py +++ b/asyncssh/connection.py @@ -179,6 +179,7 @@ _ProtocolFactory = Union[_ClientFactory, _ServerFactory] _Conn = TypeVar('_Conn', 'SSHClientConnection', 'SSHServerConnection') +_ConnSelf = TypeVar("_ConnSelf", bound="SSHConnection") class _TunnelProtocol(Protocol): """Base protocol for connections to tunnel SSH over""" @@ -893,7 +894,7 @@ def __init__(self, loop: asyncio.AbstractEventLoop, self._disable_trivial_auth = False - async def __aenter__(self) -> 'SSHConnection': + async def __aenter__(self: _ConnSelf) -> _ConnSelf: """Allow SSHConnection to be used as an async context manager""" return self From cad4aa8feecd3cf5bd9fb0574234fbf6550edd38 Mon Sep 17 00:00:00 2001 From: Pieter-Jan Briers Date: Sun, 12 Nov 2023 12:38:27 +0100 Subject: [PATCH 2/2] Use single quotes --- asyncssh/connection.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/asyncssh/connection.py b/asyncssh/connection.py index 9038ac4..2fb1ac2 100644 --- a/asyncssh/connection.py +++ b/asyncssh/connection.py @@ -179,7 +179,7 @@ _ProtocolFactory = Union[_ClientFactory, _ServerFactory] _Conn = TypeVar('_Conn', 'SSHClientConnection', 'SSHServerConnection') -_ConnSelf = TypeVar("_ConnSelf", bound="SSHConnection") +_ConnSelf = TypeVar('_ConnSelf', bound='SSHConnection') class _TunnelProtocol(Protocol): """Base protocol for connections to tunnel SSH over"""