Skip to content

Commit

Permalink
Use Self TypeVar for SSHConnection __aenter__ (#600)
Browse files Browse the repository at this point in the history
* 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.
  • Loading branch information
PJB3005 authored Nov 12, 2023
1 parent 2374e28 commit e9e95ae
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 @@ -944,7 +945,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 e9e95ae

Please sign in to comment.