From: Mike Bayer <mike_mp@zzzcomputing.com>
Date: Sat, 22 Apr 2023 13:41:49 +0000 (-0400)
Subject: support slice access for .c
X-Git-Tag: rel_2_0_11~13
X-Git-Url: http://git.ipfire.org/?a=commitdiff_plain;h=d62de424ef738a0cfb9c682d9b25a8644dff5985;p=thirdparty%2Fsqlalchemy%2Fsqlalchemy.git

support slice access for .c

Added support for slice access with :class:`.ColumnCollection`, e.g.
``table.c[0:5]``, ``subquery.c[:-1]`` etc. Slice access returns a sub
:class:`.ColumnCollection` in the same way as passing a tuple of keys. This
is a natural continuation of the key-tuple access added for :ticket:`8285`,
which it appears to be an oversight that this usage was omitted.

Change-Id: I6378642f39501ffbbae4acadf1dc38a43c39d722
References: #8285
References: #9690
---

diff --git a/doc/build/changelog/unreleased_20/c_slice.rst b/doc/build/changelog/unreleased_20/c_slice.rst
new file mode 100644
index 0000000000..c3e2e4bcc8
--- /dev/null
+++ b/doc/build/changelog/unreleased_20/c_slice.rst
@@ -0,0 +1,9 @@
+.. change::
+    :tags: usecase, sql
+    :tickets: 8285
+
+    Added support for slice access with :class:`.ColumnCollection`, e.g.
+    ``table.c[0:5]``, ``subquery.c[:-1]`` etc. Slice access returns a sub
+    :class:`.ColumnCollection` in the same way as passing a tuple of keys. This
+    is a natural continuation of the key-tuple access added for :ticket:`8285`,
+    which it appears to be an oversight that this usage was omitted.
diff --git a/lib/sqlalchemy/sql/base.py b/lib/sqlalchemy/sql/base.py
index 2539277702..309555338e 100644
--- a/lib/sqlalchemy/sql/base.py
+++ b/lib/sqlalchemy/sql/base.py
@@ -1591,14 +1591,26 @@ class ColumnCollection(Generic[_COLKEY, _COL_co]):
     ) -> ReadOnlyColumnCollection[_COLKEY, _COL_co]:
         ...
 
+    @overload
+    def __getitem__(
+        self, key: slice
+    ) -> ReadOnlyColumnCollection[_COLKEY, _COL_co]:
+        ...
+
     def __getitem__(
-        self, key: Union[str, int, Tuple[Union[str, int], ...]]
+        self, key: Union[str, int, slice, Tuple[Union[str, int], ...]]
     ) -> Union[ReadOnlyColumnCollection[_COLKEY, _COL_co], _COL_co]:
         try:
-            if isinstance(key, tuple):
-                return ColumnCollection(  # type: ignore
-                    [self._index[sub_key] for sub_key in key]
-                ).as_readonly()
+            if isinstance(key, (tuple, slice)):
+                if isinstance(key, slice):
+                    cols = (
+                        (sub_key, col)
+                        for (sub_key, col, _) in self._collection[key]
+                    )
+                else:
+                    cols = (self._index[sub_key] for sub_key in key)
+
+                return ColumnCollection(cols).as_readonly()
             else:
                 return self._index[key][1]
         except KeyError as err:
diff --git a/test/ext/mypy/plain_files/selectables.py b/test/ext/mypy/plain_files/selectables.py
new file mode 100644
index 0000000000..7d31124587
--- /dev/null
+++ b/test/ext/mypy/plain_files/selectables.py
@@ -0,0 +1,17 @@
+from sqlalchemy import column
+from sqlalchemy import table
+
+
+def test_col_accessors() -> None:
+    t = table("t", column("a"), column("b"), column("c"))
+
+    t.c.a
+    t.c["a"]
+
+    t.c[2]
+    t.c[0, 1]
+    t.c[0, 1, "b", "c"]
+    t.c[(0, 1, "b", "c")]
+
+    t.c[:-1]
+    t.c[0:2]
diff --git a/test/sql/test_select.py b/test/sql/test_select.py
index 7979fd200c..5268c41dca 100644
--- a/test/sql/test_select.py
+++ b/test/sql/test_select.py
@@ -528,6 +528,22 @@ class ColumnCollectionAsSelectTest(fixtures.TestBase, AssertsCompiledSQL):
 
         eq_(list(coll), [table1.c.description, table1.c.myid])
 
+    def test_c_sub_collection_positive_slice(self):
+        coll = table1.c[0:2]
+
+        is_(coll.myid, table1.c.myid)
+        is_(coll.name, table1.c.name)
+
+        eq_(list(coll), [table1.c.myid, table1.c.name])
+
+    def test_c_sub_collection_negative_slice(self):
+        coll = table1.c[-2:]
+
+        is_(coll.name, table1.c.name)
+        is_(coll.description, table1.c.description)
+
+        eq_(list(coll), [table1.c.name, table1.c.description])
+
     def test_missing_key(self):
 
         with expect_raises_message(KeyError, "unknown"):