From: Denis Laxalde <denis@laxalde.org>
Date: Tue, 18 Mar 2025 16:23:01 +0000 (-0400)
Subject: Make ARRAY generic on the item_type
X-Git-Url: http://git.ipfire.org/?a=commitdiff_plain;h=500adfafcb782c5b22ff49e00192a2ed42ed09b6;p=thirdparty%2Fsqlalchemy%2Fsqlalchemy.git

Make ARRAY generic on the item_type

Now `Column(type_=ARRAY(Integer)` is inferred as `Column[Sequence[int]]` instead as `Column[Sequence[Any]]` previously. This only works with the `type_` argument to Column, but that's not new.

This follows from a suggestion at
https://github.com/sqlalchemy/sqlalchemy/pull/12386#issuecomment-2694056069.

Related to #6810.

Closes: #12443
Pull-request: https://github.com/sqlalchemy/sqlalchemy/pull/12443
Pull-request-sha: 2fff4e89cd0b72d9444ce3f3d845b152770fc55d

Change-Id: I87b828fd82d10fbf157141db3c31f0ec8149caad
---

diff --git a/lib/sqlalchemy/dialects/postgresql/array.py b/lib/sqlalchemy/dialects/postgresql/array.py
index f32f146664..af026fb6ba 100644
--- a/lib/sqlalchemy/dialects/postgresql/array.py
+++ b/lib/sqlalchemy/dialects/postgresql/array.py
@@ -197,7 +197,7 @@ class array(expression.ExpressionClauseList[_T]):
             return self
 
 
-class ARRAY(sqltypes.ARRAY):
+class ARRAY(sqltypes.ARRAY[_T]):
     """PostgreSQL ARRAY type.
 
     The :class:`_postgresql.ARRAY` type is constructed in the same way
@@ -271,7 +271,7 @@ class ARRAY(sqltypes.ARRAY):
 
     def __init__(
         self,
-        item_type: _TypeEngineArgument[typing_Any],
+        item_type: _TypeEngineArgument[_T],
         as_tuple: bool = False,
         dimensions: Optional[int] = None,
         zero_indexes: bool = False,
@@ -320,7 +320,7 @@ class ARRAY(sqltypes.ARRAY):
         self.dimensions = dimensions
         self.zero_indexes = zero_indexes
 
-    class Comparator(sqltypes.ARRAY.Comparator):
+    class Comparator(sqltypes.ARRAY.Comparator[_T]):
         """Define comparison operations for :class:`_types.ARRAY`.
 
         Note that these operations are in addition to those provided
@@ -361,7 +361,7 @@ class ARRAY(sqltypes.ARRAY):
     def _against_native_enum(self) -> bool:
         return (
             isinstance(self.item_type, sqltypes.Enum)
-            and self.item_type.native_enum
+            and self.item_type.native_enum  # type: ignore[attr-defined]
         )
 
     def literal_processor(
diff --git a/lib/sqlalchemy/sql/sqltypes.py b/lib/sqlalchemy/sql/sqltypes.py
index 58af4cc0af..f71678a4ab 100644
--- a/lib/sqlalchemy/sql/sqltypes.py
+++ b/lib/sqlalchemy/sql/sqltypes.py
@@ -2813,7 +2813,7 @@ class JSON(Indexable, TypeEngine[Any]):
 
 
 class ARRAY(
-    SchemaEventTarget, Indexable, Concatenable, TypeEngine[Sequence[Any]]
+    SchemaEventTarget, Indexable, Concatenable, TypeEngine[Sequence[_T]]
 ):
     """Represent a SQL Array type.
 
@@ -2936,7 +2936,7 @@ class ARRAY(
 
     def __init__(
         self,
-        item_type: _TypeEngineArgument[Any],
+        item_type: _TypeEngineArgument[_T],
         as_tuple: bool = False,
         dimensions: Optional[int] = None,
         zero_indexes: bool = False,
@@ -2985,8 +2985,8 @@ class ARRAY(
         self.zero_indexes = zero_indexes
 
     class Comparator(
-        Indexable.Comparator[Sequence[Any]],
-        Concatenable.Comparator[Sequence[Any]],
+        Indexable.Comparator[Sequence[_T]],
+        Concatenable.Comparator[Sequence[_T]],
     ):
         """Define comparison operations for :class:`_types.ARRAY`.
 
@@ -2997,7 +2997,7 @@ class ARRAY(
 
         __slots__ = ()
 
-        type: ARRAY
+        type: ARRAY[_T]
 
         @overload
         def _setup_getitem(
diff --git a/test/typing/plain_files/dialects/postgresql/pg_stuff.py b/test/typing/plain_files/dialects/postgresql/pg_stuff.py
index 9981e4a4fc..b74ea53082 100644
--- a/test/typing/plain_files/dialects/postgresql/pg_stuff.py
+++ b/test/typing/plain_files/dialects/postgresql/pg_stuff.py
@@ -117,3 +117,9 @@ reveal_type(array_of_ints)
 
 # EXPECTED_MYPY: Cannot infer type argument 1 of "array"
 array([0], type_=Text)
+
+# EXPECTED_TYPE: ARRAY[str]
+reveal_type(ARRAY(Text))
+
+# EXPECTED_TYPE: Column[Sequence[int]]
+reveal_type(Column(type_=ARRAY(Integer)))