summaryrefslogtreecommitdiff
diff options
context:
space:
mode:
authorMike Bayer <mike_mp@zzzcomputing.com>2023-04-22 09:41:49 -0400
committerMike Bayer <mike_mp@zzzcomputing.com>2023-04-22 12:20:43 -0400
commitd62de424ef738a0cfb9c682d9b25a8644dff5985 (patch)
treef97f37ade36b3f97cb608757b9b1cbd589fe074b
parentb3e01c30d2610f5454eda5b29f093b65e54230b5 (diff)
downloadsqlalchemy-d62de424ef738a0cfb9c682d9b25a8644dff5985.tar.gz
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
-rw-r--r--doc/build/changelog/unreleased_20/c_slice.rst9
-rw-r--r--lib/sqlalchemy/sql/base.py22
-rw-r--r--test/ext/mypy/plain_files/selectables.py17
-rw-r--r--test/sql/test_select.py16
4 files changed, 59 insertions, 5 deletions
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 000000000..c3e2e4bcc
--- /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 253927770..309555338 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 000000000..7d3112458
--- /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 7979fd200..5268c41dc 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"):