Use psycopg2 type stubs (#12269)

This commit is contained in:
David Robertson 2022-03-23 14:03:24 +00:00 committed by GitHub
parent c5776780f0
commit f4c5e5864c
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
4 changed files with 14 additions and 4 deletions

1
changelog.d/12269.misc Normal file
View file

@ -0,0 +1 @@
Use type stubs for `psycopg2`.

View file

@ -108,6 +108,7 @@ CONDITIONAL_REQUIREMENTS["mypy"] = [
"types-jsonschema>=3.2.0",
"types-opentracing>=2.4.2",
"types-Pillow>=8.3.4",
"types-psycopg2>=2.9.9",
"types-pyOpenSSL>=20.0.7",
"types-PyYAML>=5.4.10",
"types-requests>=2.26.0",

View file

@ -288,7 +288,7 @@ class LoggingTransaction:
"""
if isinstance(self.database_engine, PostgresEngine):
from psycopg2.extras import execute_batch # type: ignore
from psycopg2.extras import execute_batch
self._do_execute(lambda *x: execute_batch(self.txn, *x), sql, args)
else:
@ -302,10 +302,18 @@ class LoggingTransaction:
rows (e.g. INSERTs).
"""
assert isinstance(self.database_engine, PostgresEngine)
from psycopg2.extras import execute_values # type: ignore
from psycopg2.extras import execute_values
return self._do_execute(
lambda *x: execute_values(self.txn, *x, fetch=fetch), sql, *args
# Type ignore: mypy is unhappy because if `x` is a 5-tuple, then there will
# be two values for `fetch`: one given positionally, and another given
# as a keyword argument. We might be able to fix this by
# - propagating the signature of psycopg2.extras.execute_values to this
# function, or
# - changing `*args: Any` to `values: T` for some appropriate T.
lambda *x: execute_values(self.txn, *x, fetch=fetch), # type: ignore[misc]
sql,
*args,
)
def execute(self, sql: str, *args: Any) -> None:

View file

@ -27,7 +27,7 @@ def create_engine(database_config) -> BaseDatabaseEngine:
if name == "psycopg2":
# Note that psycopg2cffi-compat provides the psycopg2 module on pypy.
import psycopg2 # type: ignore
import psycopg2
return PostgresEngine(psycopg2, database_config)