From 9fdf5d323a7a3c3d4fd9342f8712b8f5af973a7d Mon Sep 17 00:00:00 2001
From: Ali Hamdan <ali.hamdan.dev@gmail.com>
Date: Sat, 17 Feb 2024 11:22:25 +0100
Subject: [PATCH] psycoprg: make Range generic

And some other small fixes
---
 stubs/psycopg2/@tests/stubtest_allowlist.txt |  1 -
 stubs/psycopg2/psycopg2/_range.pyi           | 57 +++++++++++---------
 2 files changed, 31 insertions(+), 27 deletions(-)

diff --git a/stubs/psycopg2/@tests/stubtest_allowlist.txt b/stubs/psycopg2/@tests/stubtest_allowlist.txt
index 1d09bcabbb70..685b7d8f6f55 100644
--- a/stubs/psycopg2/@tests/stubtest_allowlist.txt
+++ b/stubs/psycopg2/@tests/stubtest_allowlist.txt
@@ -1,3 +1,2 @@
 psycopg2.pool.AbstractConnectionPool.(closeall|getconn|putconn)
-psycopg2.(extras|_range).RangeAdapter.name
 psycopg2.(_psycopg|extensions).connection.async  # async is a reserved keyword
diff --git a/stubs/psycopg2/psycopg2/_range.pyi b/stubs/psycopg2/psycopg2/_range.pyi
index 0b3247455b71..eec56727e740 100644
--- a/stubs/psycopg2/psycopg2/_range.pyi
+++ b/stubs/psycopg2/psycopg2/_range.pyi
@@ -1,17 +1,20 @@
-from _typeshed import Incomplete
-from typing import Any, overload
+import datetime as dt
+from _typeshed import SupportsAllComparisons
+from typing import Any, Generic, TypeVar, overload
 from typing_extensions import Self
 
-from psycopg2._psycopg import connection, cursor
+from psycopg2._psycopg import _type, connection, cursor
 
-class Range:
+_T_co = TypeVar("_T_co", covariant=True)
+
+class Range(Generic[_T_co]):
     def __init__(
-        self, lower: Incomplete | None = None, upper: Incomplete | None = None, bounds: str = "[)", empty: bool = False
+        self, lower: _T_co | None = None, upper: _T_co | None = None, bounds: str = "[)", empty: bool = False
     ) -> None: ...
     @property
-    def lower(self): ...
+    def lower(self) -> _T_co | None: ...
     @property
-    def upper(self): ...
+    def upper(self) -> _T_co | None: ...
     @property
     def isempty(self) -> bool: ...
     @property
@@ -22,38 +25,38 @@ class Range:
     def lower_inc(self) -> bool: ...
     @property
     def upper_inc(self) -> bool: ...
-    def __contains__(self, x) -> bool: ...
+    def __contains__(self, x: SupportsAllComparisons) -> bool: ...
     def __bool__(self) -> bool: ...
     def __eq__(self, other: object) -> bool: ...
     def __ne__(self, other: object) -> bool: ...
     def __hash__(self) -> int: ...
-    def __lt__(self, other: Range) -> bool: ...
-    def __le__(self, other: Range) -> bool: ...
-    def __gt__(self, other: Range) -> bool: ...
-    def __ge__(self, other: Range) -> bool: ...
+    def __lt__(self, other: Range[_T_co]) -> bool: ...
+    def __le__(self, other: Range[_T_co]) -> bool: ...
+    def __gt__(self, other: Range[_T_co]) -> bool: ...
+    def __ge__(self, other: Range[_T_co]) -> bool: ...
 
 def register_range(
-    pgrange: str, pyrange: str | type[Range], conn_or_curs: connection | cursor, globally: bool = False
+    pgrange: str, pyrange: str | type[Range[Any]], conn_or_curs: connection | cursor, globally: bool = False
 ) -> RangeCaster: ...
 
 class RangeAdapter:
-    name: str  # this is None here but MUST be str in subclasses
-    adapted: Range
-    def __init__(self, adapted: Range) -> None: ...
+    name: str | None = None
+    adapted: Range[Any]
+    def __init__(self, adapted: Range[Any]) -> None: ...
     def __conform__(self, proto) -> Self | None: ...
     def prepare(self, conn: connection | None) -> None: ...
     def getquoted(self) -> bytes: ...
 
 class RangeCaster:
     adapter: type[RangeAdapter]
-    range: type[Range]
-    subtype_oid: Any
-    typecaster: Any
-    array_typecaster: Any
+    range: type[Range[Any]]
+    subtype_oid: int
+    typecaster: _type
+    array_typecaster: _type | None
     def __init__(
         self,
         pgrange: str | type[RangeAdapter],
-        pyrange: str | type[Range],
+        pyrange: str | type[Range[Any]],
         oid: int,
         subtype_oid: int,
         array_oid: int | None = None,
@@ -61,12 +64,14 @@ class RangeCaster:
     @overload
     def parse(self, s: None, cur: cursor | None = None) -> None: ...
     @overload
-    def parse(self, s: str, cur: cursor | None = None) -> Range: ...
+    def parse(self, s: str, cur: cursor | None = None) -> Range[Any]: ...
+    @overload
+    def parse(self, s: str | None, cur: cursor | None = None) -> Range[Any] | None: ...
 
-class NumericRange(Range): ...
-class DateRange(Range): ...
-class DateTimeRange(Range): ...
-class DateTimeTZRange(Range): ...
+class NumericRange(Range[float]): ...
+class DateRange(Range[dt.date]): ...
+class DateTimeRange(Range[dt.datetime]): ...
+class DateTimeTZRange(Range[dt.datetime]): ...
 
 class NumberRangeAdapter(RangeAdapter):
     def getquoted(self) -> bytes: ...