From 3af2db1151491f77a2d53fb82a1d64ef3d172479 Mon Sep 17 00:00:00 2001 From: Daniele Varrazzo Date: Mon, 28 Jun 2021 03:14:20 +0100 Subject: [PATCH] Use the _psycopg module from either psycopg_c or psycopg_binary --- psycopg/psycopg/_cmodule.py | 20 +++++++++++ psycopg/psycopg/_wrappers.py | 51 +++++++++++++++++++++++++++ psycopg/psycopg/adapt.py | 10 +++--- psycopg/psycopg/connection.py | 5 ++- psycopg/psycopg/copy.py | 5 ++- psycopg/psycopg/cursor.py | 5 ++- psycopg/psycopg/types/numeric.py | 35 +++++------------- psycopg_c/psycopg_c/types/numeric.pyx | 3 +- tests/test_adapt.py | 5 ++- 9 files changed, 92 insertions(+), 47 deletions(-) create mode 100644 psycopg/psycopg/_cmodule.py create mode 100644 psycopg/psycopg/_wrappers.py diff --git a/psycopg/psycopg/_cmodule.py b/psycopg/psycopg/_cmodule.py new file mode 100644 index 000000000..0ab4813d2 --- /dev/null +++ b/psycopg/psycopg/_cmodule.py @@ -0,0 +1,20 @@ +""" +Simplify access to the _psycopg module +""" + +# Copyright (C) 2021 The Psycopg Team + +from . import pq + +# Note: "c" must the first attempt so that mypy associates the variable the +# right module interface. It will not result Optional, but hey. +if pq.__impl__ == "c": + from psycopg_c import _psycopg +elif pq.__impl__ == "binary": + from psycopg_binary import _psycopg # type: ignore +elif pq.__impl__ == "python": + _psycopg = None # type: ignore +else: + raise ImportError( + f"can't find _psycopg optimised module in {pq.__impl__!r}" + ) diff --git a/psycopg/psycopg/_wrappers.py b/psycopg/psycopg/_wrappers.py new file mode 100644 index 000000000..ecbb34c3a --- /dev/null +++ b/psycopg/psycopg/_wrappers.py @@ -0,0 +1,51 @@ +""" +Wrappers for numeric types. +""" + +# Copyright (C) 2020-2021 The Psycopg Team + +# Wrappers to force numbers to be cast as specific PostgreSQL types + +# These types are implemented here but exposed by `psycopg.types.numeric`. +# They are defined here to avoid a circular import. +_MODULE = "psycopg.types.numeric" + + +class Int2(int): + + __module__ = _MODULE + + def __new__(cls, arg: int) -> "Int2": + return super().__new__(cls, arg) + + +class Int4(int): + + __module__ = _MODULE + + def __new__(cls, arg: int) -> "Int4": + return super().__new__(cls, arg) + + +class Int8(int): + + __module__ = _MODULE + + def __new__(cls, arg: int) -> "Int8": + return super().__new__(cls, arg) + + +class IntNumeric(int): + + __module__ = _MODULE + + def __new__(cls, arg: int) -> "IntNumeric": + return super().__new__(cls, arg) + + +class Oid(int): + + __module__ = _MODULE + + def __new__(cls, arg: int) -> "Oid": + return super().__new__(cls, arg) diff --git a/psycopg/psycopg/adapt.py b/psycopg/psycopg/adapt.py index 06b404ba8..6acc9bb1f 100644 --- a/psycopg/psycopg/adapt.py +++ b/psycopg/psycopg/adapt.py @@ -14,6 +14,7 @@ from . import errors as e from ._enums import Format as Format from .oids import postgres_types from .proto import AdaptContext, Buffer as Buffer +from ._cmodule import _psycopg from ._typeinfo import TypesRegistry if TYPE_CHECKING: @@ -186,7 +187,7 @@ class AdaptersMap(AdaptContext): f"dumpers should be registered on classes, got {cls} instead" ) - if pq.__impl__ != "python": + if _psycopg: dumper = self._get_optimised(dumper) # Register the dumper both as its format and as default @@ -210,7 +211,7 @@ class AdaptersMap(AdaptContext): f"loaders should be registered on oid, got {oid} instead" ) - if pq.__impl__ != "python": + if _psycopg: loader = self._get_optimised(loader) fmt = loader.format @@ -272,7 +273,6 @@ class AdaptersMap(AdaptContext): # Check if the class comes from psycopg.types and there is a class # with the same name in psycopg_c._psycopg. from psycopg import types - from psycopg_c import _psycopg if cls.__module__.startswith(types.__name__): new = cast(Type[RV], getattr(_psycopg, cls.__name__, None)) @@ -292,9 +292,7 @@ global_adapters = AdaptersMap(types=postgres_types) Transformer: Type[proto.Transformer] # Override it with fast object if available -if pq.__impl__ == "c": - from psycopg_c import _psycopg - +if _psycopg: Transformer = _psycopg.Transformer else: from . import _transform diff --git a/psycopg/psycopg/connection.py b/psycopg/psycopg/connection.py index 906b5fdeb..5faedf261 100644 --- a/psycopg/psycopg/connection.py +++ b/psycopg/psycopg/connection.py @@ -28,6 +28,7 @@ from .proto import AdaptContext, ConnectionType, Params, PQGen, PQGenConn from .proto import Query, RV from .compat import asynccontextmanager from .cursor import Cursor, AsyncCursor +from ._cmodule import _psycopg from .conninfo import _conninfo_connect_timeout, ConnectionInfo from .generators import notifies from ._preparing import PrepareManager @@ -47,9 +48,7 @@ if TYPE_CHECKING: from .pq.proto import PGconn, PGresult from .pool.base import BasePool -if pq.__impl__ == "c": - from psycopg_c import _psycopg - +if _psycopg: connect = _psycopg.connect execute = _psycopg.execute diff --git a/psycopg/psycopg/copy.py b/psycopg/psycopg/copy.py index f34868ede..f258be088 100644 --- a/psycopg/psycopg/copy.py +++ b/psycopg/psycopg/copy.py @@ -20,6 +20,7 @@ from .pq import ExecStatus from .adapt import Format from .proto import ConnectionType, PQGen, Transformer from .compat import create_task +from ._cmodule import _psycopg from .generators import copy_from, copy_to, copy_end if TYPE_CHECKING: @@ -639,9 +640,7 @@ def _load_sub( # Override functions with fast versions if available -if pq.__impl__ == "c": - from psycopg_c import _psycopg - +if _psycopg: format_row_text = _psycopg.format_row_text format_row_binary = _psycopg.format_row_binary parse_row_text = _psycopg.parse_row_text diff --git a/psycopg/psycopg/cursor.py b/psycopg/psycopg/cursor.py index 0fd8a0573..3fa4ac93c 100644 --- a/psycopg/psycopg/cursor.py +++ b/psycopg/psycopg/cursor.py @@ -21,6 +21,7 @@ from .rows import Row, RowFactory from .proto import ConnectionType, Query, Params, PQGen from .compat import asynccontextmanager from ._column import Column +from ._cmodule import _psycopg from ._queries import PostgresQuery from ._preparing import Prepare @@ -32,9 +33,7 @@ if TYPE_CHECKING: execute: Callable[["PGconn"], PQGen[List["PGresult"]]] -if pq.__impl__ == "c": - from psycopg_c import _psycopg - +if _psycopg: execute = _psycopg.execute else: diff --git a/psycopg/psycopg/types/numeric.py b/psycopg/psycopg/types/numeric.py index decfa2320..2d12af3fe 100644 --- a/psycopg/psycopg/types/numeric.py +++ b/psycopg/psycopg/types/numeric.py @@ -20,33 +20,14 @@ from .._struct import pack_int4, pack_uint4, unpack_int4, unpack_uint4 from .._struct import pack_int8, unpack_int8 from .._struct import pack_float8, unpack_float4, unpack_float8 - -# Wrappers to force numbers to be cast as specific PostgreSQL types - - -class Int2(int): - def __new__(cls, arg: int) -> "Int2": - return super().__new__(cls, arg) - - -class Int4(int): - def __new__(cls, arg: int) -> "Int4": - return super().__new__(cls, arg) - - -class Int8(int): - def __new__(cls, arg: int) -> "Int8": - return super().__new__(cls, arg) - - -class IntNumeric(int): - def __new__(cls, arg: int) -> "IntNumeric": - return super().__new__(cls, arg) - - -class Oid(int): - def __new__(cls, arg: int) -> "Oid": - return super().__new__(cls, arg) +# Exposed here +from .._wrappers import ( + Int2 as Int2, + Int4 as Int4, + Int8 as Int8, + IntNumeric as IntNumeric, + Oid as Oid, +) class _NumberDumper(Dumper): diff --git a/psycopg_c/psycopg_c/types/numeric.pyx b/psycopg_c/psycopg_c/types/numeric.pyx index 01a67b3a3..75af361b1 100644 --- a/psycopg_c/psycopg_c/types/numeric.pyx +++ b/psycopg_c/psycopg_c/types/numeric.pyx @@ -21,8 +21,7 @@ from decimal import Decimal, Context, DefaultContext from psycopg_c._psycopg cimport endian from psycopg import errors as e - -from psycopg.types.numeric import Int2, Int4, Int8, IntNumeric +from psycopg._wrappers import Int2, Int4, Int8, IntNumeric cdef extern from "Python.h": # work around https://github.com/cython/cython/issues/3909 diff --git a/tests/test_adapt.py b/tests/test_adapt.py index 3988d7d37..6815d27f1 100644 --- a/tests/test_adapt.py +++ b/tests/test_adapt.py @@ -7,6 +7,7 @@ import psycopg from psycopg import pq from psycopg.adapt import Transformer, Format, Dumper, Loader from psycopg.oids import postgres_types as builtins, TEXT_OID +from psycopg._cmodule import _psycopg @pytest.mark.parametrize( @@ -292,11 +293,9 @@ def test_no_cast_needed(conn, fmt_in): assert cur.fetchone()[0] == 20 -@pytest.mark.skipif(psycopg.pq.__impl__ == "python", reason="C module test") +@pytest.mark.skipif(_psycopg is None, reason="C module test") def test_optimised_adapters(): - from psycopg_c import _psycopg - # All the optimised adapters available c_adapters = {} for n in dir(_psycopg): -- 2.47.3