Skip to content

Commit

Permalink
Use Self TypeVar for SSHConnection __aenter__
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
PJB3005 committed Oct 17, 2023
1 parent 5ca8487 commit ea46520
Showing 1 changed file with 2 additions and 1 deletion.
3 changes: 2 additions & 1 deletion asyncssh/connection.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"""
Expand Down Expand Up @@ -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
Expand Down

0 comments on commit ea46520

Please sign in to comment.