Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add typing to two objects in connection_utils #1198

Merged
merged 3 commits into from
Oct 29, 2024
Merged

Add typing to two objects in connection_utils #1198

Changes from 2 commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
a31d182
Add type annotations to `TLSUpgradeProto`
DanielNoord Oct 22, 2024
76105cc
Add type annotations to `_create_ssl_connection`
DanielNoord Oct 23, 2024
6905497
Import `__future__` and unstringify annotation
DanielNoord Oct 24, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
34 changes: 29 additions & 5 deletions asyncpg/connect_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@

import asyncio
import collections
from collections.abc import Callable
import enum
import functools
import getpass
Expand Down Expand Up @@ -764,14 +765,21 @@ def _parse_connect_arguments(*, dsn, host, port, user, password, passfile,


class TLSUpgradeProto(asyncio.Protocol):
def __init__(self, loop, host, port, ssl_context, ssl_is_advisory):
def __init__(
self,
loop: asyncio.AbstractEventLoop,
host: str,
port: int,
ssl_context: ssl_module.SSLContext,
ssl_is_advisory: bool,
) -> None:
self.on_data = _create_future(loop)
self.host = host
self.port = port
self.ssl_context = ssl_context
self.ssl_is_advisory = ssl_is_advisory

def data_received(self, data):
def data_received(self, data: bytes) -> None:
if data == b'S':
self.on_data.set_result(True)
elif (self.ssl_is_advisory and
Expand All @@ -789,15 +797,30 @@ def data_received(self, data):
'rejected SSL upgrade'.format(
host=self.host, port=self.port)))

def connection_lost(self, exc):
def connection_lost(self, exc: typing.Optional[Exception]) -> None:
if not self.on_data.done():
if exc is None:
exc = ConnectionError('unexpected connection_lost() call')
self.on_data.set_exception(exc)


async def _create_ssl_connection(protocol_factory, host, port, *,
loop, ssl_context, ssl_is_advisory=False):
_ProctolFactoryR = typing.TypeVar(
"_ProctolFactoryR", bound=asyncio.protocols.Protocol
)


async def _create_ssl_connection(
# TODO: The return type is a specific combination of subclasses of
# asyncio.protocols.Protocol that we can't express. For now, having the
# return type be dependent on signature of the factory is an improvement
protocol_factory: "Callable[[], _ProctolFactoryR]",
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
protocol_factory: "Callable[[], _ProctolFactoryR]",
protocol_factory: Callable[[], _ProctolFactoryR],

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

That's not possible sadly. Callable isn't subscribtable at runtime like this on Python 3.8. See earlier CI runs in which CI complained about this.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

For compatibility you want typing.Callable, not collections.abc.Callable

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In later versions collections is preferred though. Isn't the use of stringified annotations the way to get compatibility?

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If you add from __future__ import annotations to the top of the file, you won't have to use strings in the annotation (and importing from collections.abc will also work)

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Also fine with me, @elprans what has your preference?

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah, do that please.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed!

host: str,
port: int,
*,
loop: asyncio.AbstractEventLoop,
ssl_context: ssl_module.SSLContext,
ssl_is_advisory: bool = False,
) -> typing.Tuple[asyncio.Transport, _ProctolFactoryR]:

tr, pr = await loop.create_connection(
lambda: TLSUpgradeProto(loop, host, port,
Expand All @@ -817,6 +840,7 @@ async def _create_ssl_connection(protocol_factory, host, port, *,
try:
new_tr = await loop.start_tls(
tr, pr, ssl_context, server_hostname=host)
assert new_tr is not None
except (Exception, asyncio.CancelledError):
tr.close()
raise
Expand Down
Loading