From 47658dbb4f8994c5b2566840632921ce8d66cb36 Mon Sep 17 00:00:00 2001 From: David Robertson Date: Thu, 12 May 2022 22:26:56 +0100 Subject: [PATCH 1/5] Annotations for psycopg2.ConnectionInfo These annotations come from the documentation here: https://www.psycopg.org/docs/extensions.html#psycopg2.extensions.ConnectionInfo If there was doubt, I referred to the libpq documentation cited by psycopg2's docs. I wasn't completely sure about `dsn_parameters`. Psycopg2's docs list it as an `dict`, and the example suggests it's a `dict[str, str]` at that. From psycopg2's source I found https://github.com/psycopg/psycopg2/blob/1d3a89a0bba621dc1cc9b32db6d241bd2da85ad1/psycopg/conninfo_type.c#L183-L206 which is implemented here: https://github.com/psycopg/psycopg2/blob/1d3a89a0bba621dc1cc9b32db6d241bd2da85ad1/psycopg/utils.c#L251-L279 I'm no expert in CPython's API, but this looks to me like it's building a `dict[str, str]`. Additionally, the libpq docs https://www.postgresql.org/docs/current/libpq-connect.html#LIBPQ-PQCONNINFO https://www.postgresql.org/docs/current/libpq-connect.html#LIBPQ-PQCONNDEFAULTS show that the underlying data just consists of strings. Additionally, I'm pretty sure from this chunk of source https://github.com/psycopg/psycopg2/blob/1d3a89a0bba621dc1cc9b32db6d241bd2da85ad1/psycopg/conninfo_type.c#L581-L598 That `ConnectionInfo.__init__` takes one positional-only argument, which must be a `psycopg2.connection`. But I don't think users are intended to be constructing this type, so I've not added that annotation. --- stubs/psycopg2/psycopg2/_psycopg.pyi | 40 ++++++++++++++-------------- 1 file changed, 20 insertions(+), 20 deletions(-) diff --git a/stubs/psycopg2/psycopg2/_psycopg.pyi b/stubs/psycopg2/psycopg2/_psycopg.pyi index 9e59da03bf81..6c0b35ecddca 100644 --- a/stubs/psycopg2/psycopg2/_psycopg.pyi +++ b/stubs/psycopg2/psycopg2/_psycopg.pyi @@ -157,27 +157,27 @@ class Column: def __setstate__(self, state): ... class ConnectionInfo: - backend_pid: Any - dbname: Any - dsn_parameters: Any - error_message: Any - host: Any - needs_password: Any - options: Any - password: Any - port: Any - protocol_version: Any - server_version: Any - socket: Any - ssl_attribute_names: Any - ssl_in_use: Any - status: Any - transaction_status: Any - used_password: Any - user: Any + backend_pid: int + dbname: str + dsn_parameters: dict[str, str] + error_message: str | None + host: str + needs_password: bool + options: str + password: str + port: int + protocol_version: int + server_version: int + socket: int + ssl_attribute_names: list[str] + ssl_in_use: bool + status: int + transaction_status: int + used_password: bool + user: str def __init__(self, *args, **kwargs) -> None: ... - def parameter_status(self, *args, **kwargs): ... - def ssl_attribute(self, *args, **kwargs): ... + def parameter_status(self, name: str) -> str | None: ... + def ssl_attribute(self, name: str) -> str | None: ... class DataError(psycopg2.DatabaseError): ... class DatabaseError(psycopg2.Error): ... From 8a0efb8129ca03d376e1af3e4b9c3b50404e5567 Mon Sep 17 00:00:00 2001 From: David Robertson Date: Thu, 12 May 2022 23:13:34 +0100 Subject: [PATCH 2/5] Annotate `connection.info` and related attributes --- stubs/psycopg2/psycopg2/_psycopg.pyi | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/stubs/psycopg2/psycopg2/_psycopg.pyi b/stubs/psycopg2/psycopg2/_psycopg.pyi index 6c0b35ecddca..4063d7774c7a 100644 --- a/stubs/psycopg2/psycopg2/_psycopg.pyi +++ b/stubs/psycopg2/psycopg2/_psycopg.pyi @@ -342,14 +342,14 @@ class connection: deferrable: Any dsn: Any encoding: Any - info: Any + info: ConnectionInfo isolation_level: Any notices: Any notifies: Any pgconn_ptr: Any - protocol_version: Any + protocol_version: int readonly: Any - server_version: Any + server_version: int status: Any string_types: Any def __init__(self, *args, **kwargs) -> None: ... @@ -361,11 +361,11 @@ class connection: @overload def cursor(self, name=..., cursor_factory: Callable[..., _T_cur] = ..., scrollable=..., withhold=...) -> _T_cur: ... def fileno(self, *args, **kwargs): ... - def get_backend_pid(self, *args, **kwargs): ... - def get_dsn_parameters(self, *args, **kwargs): ... + def get_backend_pid(self) -> int: ... + def get_dsn_parameters(self) -> dict[str, str]: ... def get_native_connection(self, *args, **kwargs): ... - def get_parameter_status(self, parameter): ... - def get_transaction_status(self): ... + def get_parameter_status(self, parameter: str) -> str | None: ... + def get_transaction_status(self) -> int: ... def isexecuting(self, *args, **kwargs): ... def lobject(self, oid=..., mode=..., new_oid=..., new_file=..., lobject_factory=...): ... def poll(self, *args, **kwargs): ... From aea1c8d7c060f19ec9f136307b73844ad86d0c80 Mon Sep 17 00:00:00 2001 From: David Robertson Date: Fri, 20 May 2022 13:54:59 +0100 Subject: [PATCH 3/5] Make ConnectionInfo attributes properties According to https://github.com/psycopg/psycopg2/blob/1d3a89a0bba621dc1cc9b32db6d241bd2da85ad1/psycopg/conninfo_type.c#L534-L563 --- stubs/psycopg2/psycopg2/_psycopg.pyi | 54 ++++++++++++++++++---------- 1 file changed, 36 insertions(+), 18 deletions(-) diff --git a/stubs/psycopg2/psycopg2/_psycopg.pyi b/stubs/psycopg2/psycopg2/_psycopg.pyi index 4063d7774c7a..b67c73c6acbc 100644 --- a/stubs/psycopg2/psycopg2/_psycopg.pyi +++ b/stubs/psycopg2/psycopg2/_psycopg.pyi @@ -157,24 +157,42 @@ class Column: def __setstate__(self, state): ... class ConnectionInfo: - backend_pid: int - dbname: str - dsn_parameters: dict[str, str] - error_message: str | None - host: str - needs_password: bool - options: str - password: str - port: int - protocol_version: int - server_version: int - socket: int - ssl_attribute_names: list[str] - ssl_in_use: bool - status: int - transaction_status: int - used_password: bool - user: str + @property + def backend_pid(self) -> int: ... + @property + def dbname(self) -> str: ... + @property + def dsn_parameters(self) -> dict[str, str]: ... + @property + def error_message(self) -> str | None: ... + @property + def host(self) -> str: ... + @property + def needs_password(self) -> bool: ... + @property + def options(self) -> str: ... + @property + def password(self) -> str: ... + @property + def port(self) -> int: ... + @property + def protocol_version(self) -> int: ... + @property + def server_version(self) -> int: ... + @property + def socket(self) -> int: ... + @property + def ssl_attribute_names(self) -> list[str]: ... + @property + def ssl_in_use(self) -> bool: ... + @property + def status(self) -> int: ... + @property + def transaction_status(self) -> int: ... + @property + def used_password(self) -> bool: ... + @property + def user(self) -> str: ... def __init__(self, *args, **kwargs) -> None: ... def parameter_status(self, name: str) -> str | None: ... def ssl_attribute(self, name: str) -> str | None: ... From 60389ae60c9f43ec87a5166b701ca0e5914a95c2 Mon Sep 17 00:00:00 2001 From: David Robertson Date: Fri, 20 May 2022 14:00:47 +0100 Subject: [PATCH 4/5] Mark connection attributes as readonly according to https://github.com/psycopg/psycopg2/blob/8ef195f2ff187454cc709d7857235676bb4176ee/psycopg/connection_type.c#L1244 --- stubs/psycopg2/psycopg2/_psycopg.pyi | 33 ++++++++++++++++++---------- 1 file changed, 22 insertions(+), 11 deletions(-) diff --git a/stubs/psycopg2/psycopg2/_psycopg.pyi b/stubs/psycopg2/psycopg2/_psycopg.pyi index b67c73c6acbc..6ee18b3f962d 100644 --- a/stubs/psycopg2/psycopg2/_psycopg.pyi +++ b/stubs/psycopg2/psycopg2/_psycopg.pyi @@ -352,24 +352,35 @@ class connection: OperationalError: Any ProgrammingError: Any Warning: Any - async_: Any + @property + def async_(self) -> Any: ... autocommit: Any - binary_types: Any - closed: Any + @property + def binary_types(self) -> Any: ... + @property + def closed(self) -> Any: ... cursor_factory: Callable[..., _cursor] deferrable: Any - dsn: Any - encoding: Any - info: ConnectionInfo + @property + def dsn(self) -> Any: ... + @property + def encoding(self) -> Any: ... + @property + def info(self) -> ConnectionInfo: ... isolation_level: Any notices: Any notifies: Any - pgconn_ptr: Any - protocol_version: int + @property + def pgconn_ptr(self) -> Any: ... + @property + def protocol_version(self) -> int: ... readonly: Any - server_version: int - status: Any - string_types: Any + @property + def server_version(self) -> int: ... + @property + def status(self) -> Any: ... + @property + def string_types(self) -> Any: ... def __init__(self, *args, **kwargs) -> None: ... def cancel(self, *args, **kwargs): ... def close(self, *args, **kwargs): ... From ed4d8df3d370231ae963aea64af9b4b6427abf12 Mon Sep 17 00:00:00 2001 From: David Robertson Date: Fri, 20 May 2022 18:04:38 +0100 Subject: [PATCH 5/5] Explain why some properties aren't `T | None` --- stubs/psycopg2/psycopg2/_psycopg.pyi | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/stubs/psycopg2/psycopg2/_psycopg.pyi b/stubs/psycopg2/psycopg2/_psycopg.pyi index 6ee18b3f962d..5c130ae8f4aa 100644 --- a/stubs/psycopg2/psycopg2/_psycopg.pyi +++ b/stubs/psycopg2/psycopg2/_psycopg.pyi @@ -157,6 +157,24 @@ class Column: def __setstate__(self, state): ... class ConnectionInfo: + # Note: the following properties can be None if their corresponding libpq function + # returns NULL. They're not annotated as such, because this is very unlikely in + # practice---the psycopg2 docs [1] don't even mention this as a possibility! + # + # - db_name + # - user + # - password + # - host + # - port + # - options + # + # (To prove this, one needs to inspect the psycopg2 source code [2], plus the + # documentation [3] and source code [4] of the corresponding libpq calls.) + # + # [1]: https://www.psycopg.org/docs/extensions.html#psycopg2.extensions.ConnectionInfo + # [2]: https://github.com/psycopg/psycopg2/blob/1d3a89a0bba621dc1cc9b32db6d241bd2da85ad1/psycopg/conninfo_type.c#L52 and below + # [3]: https://www.postgresql.org/docs/current/libpq-status.html + # [4]: https://github.com/postgres/postgres/blob/b39838889e76274b107935fa8e8951baf0e8b31b/src/interfaces/libpq/fe-connect.c#L6754 and below @property def backend_pid(self) -> int: ... @property