refactor: less verbose pattern to manage optional libpq functions

Author: dvarrazzo

psycopg/psycopg/pq/_pq_ctypes.py

@@ -9,7 +9,7 @@
 import ctypes.util
 from ctypes import Structure, CFUNCTYPE, POINTER
 from ctypes import c_char, c_char_p, c_int, c_size_t, c_ubyte, c_uint, c_void_p
-from typing import Any, List, Optional, Tuple
+from typing import Any, List, NoReturn, Tuple
 
 from .misc import find_libpq_full_path
 from ..errors import NotSupportedError
@@ -185,24 +185,23 @@ class PGresAttDesc_struct(Structure):
 PQhost.argtypes = [PGconn_ptr]
 PQhost.restype = c_char_p
 
-_PQhostaddr = None
 
-if libpq_version >= 120000:
-    _PQhostaddr = pq.PQhostaddr
-    _PQhostaddr.argtypes = [PGconn_ptr]
-    _PQhostaddr.restype = c_char_p
-
-
-def PQhostaddr(pgconn: PGconn_struct) -> bytes:
-    if not _PQhostaddr:
+def not_supported_before(fname: str, pgversion: int) -> Any:
+    def not_supported(*args: Any, **kwargs: Any) -> NoReturn:
         raise NotSupportedError(
-            "PQhostaddr requires libpq from PostgreSQL 12,"
-            f" {libpq_version} available instead"
+            f"{fname} requires libpq from PostgreSQL {pgversion} on the client;"
+            f" version {libpq_version // 10000} available instead"
         )
 
-    rv: bytes = _PQhostaddr(pgconn)
-    return rv
+    return not_supported
+
 
+if libpq_version >= 120000:
+    PQhostaddr = pq.PQhostaddr
+    PQhostaddr.argtypes = [PGconn_ptr]
+    PQhostaddr.restype = c_char_p
+else:
+    PQhostaddr = not_supported_before("PQhostaddr", 12)
 
 PQport = pq.PQport
 PQport.argtypes = [PGconn_ptr]
@@ -306,38 +305,18 @@ def PQhostaddr(pgconn: PGconn_struct) -> bytes:
 PQdescribePortal.argtypes = [PGconn_ptr, c_char_p]
 PQdescribePortal.restype = PGresult_ptr
 
-_PQclosePrepared = None
-_PQclosePortal = None
-
 if libpq_version >= 170000:
-    _PQclosePrepared = pq.PQclosePrepared
-    _PQclosePrepared.argtypes = [PGconn_ptr, c_char_p]
-    _PQclosePrepared.restype = PGresult_ptr
-
-    _PQclosePortal = pq.PQclosePortal
-    _PQclosePortal.argtypes = [PGconn_ptr, c_char_p]
-    _PQclosePortal.restype = PGresult_ptr
-
+    PQclosePrepared = pq.PQclosePrepared
+    PQclosePrepared.argtypes = [PGconn_ptr, c_char_p]
+    PQclosePrepared.restype = PGresult_ptr
 
-def PQclosePrepared(pgconn: PGconn_struct, name: str) -> int:
-    if not _PQclosePrepared:
-        raise NotSupportedError(
-            "PQclosePrepared requires libpq from PostgreSQL 17,"
-            f" {libpq_version} available instead"
-        )
-    rv: int = _PQclosePrepared(pgconn, name)
-    return rv
-
-
-def PQclosePortal(pgconn: PGconn_struct, name: str) -> int:
-    if not _PQclosePortal:
-        raise NotSupportedError(
-            "PQclosePortal requires libpq from PostgreSQL 17,"
-            f" {libpq_version} available instead"
-        )
-    rv: int = _PQclosePortal(pgconn, name)
-    return rv
+    PQclosePortal = pq.PQclosePortal
+    PQclosePortal.argtypes = [PGconn_ptr, c_char_p]
+    PQclosePortal.restype = PGresult_ptr
 
+else:
+    PQclosePrepared = not_supported_before("PQclosePrepared", 17)
+    PQclosePortal = not_supported_before("PQclosePrepared", 17)
 
 PQresultStatus = pq.PQresultStatus
 PQresultStatus.argtypes = [PGresult_ptr]
@@ -532,38 +511,18 @@ def PQclosePortal(pgconn: PGconn_struct, name: str) -> int:
 PQsendDescribePortal.argtypes = [PGconn_ptr, c_char_p]
 PQsendDescribePortal.restype = c_int
 
-_PQsendClosePrepared = None
-_PQsendClosePortal = None
-
 if libpq_version >= 170000:
-    _PQsendClosePrepared = pq.PQsendClosePrepared
-    _PQsendClosePrepared.argtypes = [PGconn_ptr, c_char_p]
-    _PQsendClosePrepared.restype = c_int
+    PQsendClosePrepared = pq.PQsendClosePrepared
+    PQsendClosePrepared.argtypes = [PGconn_ptr, c_char_p]
+    PQsendClosePrepared.restype = c_int
 
-    _PQsendClosePortal = pq.PQsendClosePortal
-    _PQsendClosePortal.argtypes = [PGconn_ptr, c_char_p]
-    _PQsendClosePortal.restype = c_int
-
-
-def PQsendClosePrepared(pgconn: PGconn_struct, name: str) -> int:
-    if not _PQsendClosePrepared:
-        raise NotSupportedError(
-            "PQsendClosePrepared requires libpq from PostgreSQL 17,"
-            f" {libpq_version} available instead"
-        )
-    rv: int = _PQsendClosePrepared(pgconn, name)
-    return rv
-
-
-def PQsendClosePortal(pgconn: PGconn_struct, name: str) -> int:
-    if not _PQsendClosePortal:
-        raise NotSupportedError(
-            "PQsendClosePortal requires libpq from PostgreSQL 17,"
-            f" {libpq_version} available instead"
-        )
-    rv: int = _PQsendClosePortal(pgconn, name)
-    return rv
+    PQsendClosePortal = pq.PQsendClosePortal
+    PQsendClosePortal.argtypes = [PGconn_ptr, c_char_p]
+    PQsendClosePortal.restype = c_int
 
+else:
+    PQsendClosePrepared = not_supported_before("PQsendClosePrepared", 17)
+    PQsendClosePortal = not_supported_before("PQsendClosePortal", 17)
 
 PQgetResult = pq.PQgetResult
 PQgetResult.argtypes = [PGconn_ptr]
@@ -640,23 +599,12 @@ def PQsendClosePortal(pgconn: PGconn_struct, name: str) -> int:
 PQtrace.argtypes = [PGconn_ptr, FILE_ptr]
 PQtrace.restype = None
 
-_PQsetTraceFlags = None
-
 if libpq_version >= 140000:
-    _PQsetTraceFlags = pq.PQsetTraceFlags
-    _PQsetTraceFlags.argtypes = [PGconn_ptr, c_int]
-    _PQsetTraceFlags.restype = None
-
-
-def PQsetTraceFlags(pgconn: PGconn_struct, flags: int) -> None:
-    if not _PQsetTraceFlags:
-        raise NotSupportedError(
-            "PQsetTraceFlags requires libpq from PostgreSQL 14,"
-            f" {libpq_version} available instead"
-        )
-
-    _PQsetTraceFlags(pgconn, flags)
-
+    PQsetTraceFlags = pq.PQsetTraceFlags
+    PQsetTraceFlags.argtypes = [PGconn_ptr, c_int]
+    PQsetTraceFlags.restype = None
+else:
+    PQsetTraceFlags = not_supported_before("PQsetTraceFlags", 14)
 
 PQuntrace = pq.PQuntrace
 PQuntrace.argtypes = [PGconn_ptr]
@@ -669,28 +617,16 @@ def PQsetTraceFlags(pgconn: PGconn_struct, flags: int) -> None:
 PQfreemem.restype = None
 
 if libpq_version >= 100000:
-    _PQencryptPasswordConn = pq.PQencryptPasswordConn
-    _PQencryptPasswordConn.argtypes = [
+    PQencryptPasswordConn = pq.PQencryptPasswordConn
+    PQencryptPasswordConn.argtypes = [
         PGconn_ptr,
         c_char_p,
         c_char_p,
         c_char_p,
     ]
-    _PQencryptPasswordConn.restype = POINTER(c_char)
-
-
-def PQencryptPasswordConn(
-    pgconn: PGconn_struct, passwd: bytes, user: bytes, algorithm: bytes
-) -> Optional[bytes]:
-    if not _PQencryptPasswordConn:
-        raise NotSupportedError(
-            "PQencryptPasswordConn requires libpq from PostgreSQL 10,"
-            f" {libpq_version} available instead"
-        )
-
-    rv: Optional[bytes] = _PQencryptPasswordConn(pgconn, passwd, user, algorithm)
-    return rv
-
+    PQencryptPasswordConn.restype = POINTER(c_char)
+else:
+    PQencryptPasswordConn = not_supported_before("PQencryptPasswordConn", 10)
 
 PQmakeEmptyPGresult = pq.PQmakeEmptyPGresult
 PQmakeEmptyPGresult.argtypes = [PGconn_ptr, c_int]
@@ -711,82 +647,33 @@ def PQencryptPasswordConn(
 
 # 34.5 Pipeline Mode
 
-_PQpipelineStatus = None
-_PQenterPipelineMode = None
-_PQexitPipelineMode = None
-_PQpipelineSync = None
-_PQsendFlushRequest = None
-
 if libpq_version >= 140000:
-    _PQpipelineStatus = pq.PQpipelineStatus
-    _PQpipelineStatus.argtypes = [PGconn_ptr]
-    _PQpipelineStatus.restype = c_int
-
-    _PQenterPipelineMode = pq.PQenterPipelineMode
-    _PQenterPipelineMode.argtypes = [PGconn_ptr]
-    _PQenterPipelineMode.restype = c_int
-
-    _PQexitPipelineMode = pq.PQexitPipelineMode
-    _PQexitPipelineMode.argtypes = [PGconn_ptr]
-    _PQexitPipelineMode.restype = c_int
-
-    _PQpipelineSync = pq.PQpipelineSync
-    _PQpipelineSync.argtypes = [PGconn_ptr]
-    _PQpipelineSync.restype = c_int
-
-    _PQsendFlushRequest = pq.PQsendFlushRequest
-    _PQsendFlushRequest.argtypes = [PGconn_ptr]
-    _PQsendFlushRequest.restype = c_int
-
-
-def PQpipelineStatus(pgconn: PGconn_struct) -> int:
-    if not _PQpipelineStatus:
-        raise NotSupportedError(
-            "PQpipelineStatus requires libpq from PostgreSQL 14,"
-            f" {libpq_version} available instead"
-        )
-    rv: int = _PQpipelineStatus(pgconn)
-    return rv
-
-
-def PQenterPipelineMode(pgconn: PGconn_struct) -> int:
-    if not _PQenterPipelineMode:
-        raise NotSupportedError(
-            "PQenterPipelineMode requires libpq from PostgreSQL 14,"
-            f" {libpq_version} available instead"
-        )
-    rv: int = _PQenterPipelineMode(pgconn)
-    return rv
-
-
-def PQexitPipelineMode(pgconn: PGconn_struct) -> int:
-    if not _PQexitPipelineMode:
-        raise NotSupportedError(
-            "PQexitPipelineMode requires libpq from PostgreSQL 14,"
-            f" {libpq_version} available instead"
-        )
-    rv: int = _PQexitPipelineMode(pgconn)
-    return rv
-
-
-def PQpipelineSync(pgconn: PGconn_struct) -> int:
-    if not _PQpipelineSync:
-        raise NotSupportedError(
-            "PQpipelineSync requires libpq from PostgreSQL 14,"
-            f" {libpq_version} available instead"
-        )
-    rv: int = _PQpipelineSync(pgconn)
-    return rv
-
-
-def PQsendFlushRequest(pgconn: PGconn_struct) -> int:
-    if not _PQsendFlushRequest:
-        raise NotSupportedError(
-            "PQsendFlushRequest requires libpq from PostgreSQL 14,"
-            f" {libpq_version} available instead"
-        )
-    rv: int = _PQsendFlushRequest(pgconn)
-    return rv
+    PQpipelineStatus = pq.PQpipelineStatus
+    PQpipelineStatus.argtypes = [PGconn_ptr]
+    PQpipelineStatus.restype = c_int
+
+    PQenterPipelineMode = pq.PQenterPipelineMode
+    PQenterPipelineMode.argtypes = [PGconn_ptr]
+    PQenterPipelineMode.restype = c_int
+
+    PQexitPipelineMode = pq.PQexitPipelineMode
+    PQexitPipelineMode.argtypes = [PGconn_ptr]
+    PQexitPipelineMode.restype = c_int
+
+    PQpipelineSync = pq.PQpipelineSync
+    PQpipelineSync.argtypes = [PGconn_ptr]
+    PQpipelineSync.restype = c_int
+
+    PQsendFlushRequest = pq.PQsendFlushRequest
+    PQsendFlushRequest.argtypes = [PGconn_ptr]
+    PQsendFlushRequest.restype = c_int
+
+else:
+    PQpipelineStatus = not_supported_before("PQpipelineStatus", 14)
+    PQenterPipelineMode = not_supported_before("PQenterPipelineMode", 14)
+    PQexitPipelineMode = not_supported_before("PQexitPipelineMode", 14)
+    PQpipelineSync = not_supported_before("PQpipelineSync", 14)
+    PQsendFlushRequest = not_supported_before("PQsendFlushRequest", 14)
 
 
 # 33.18. SSL Support