diff options
author | Mike Bayer <mike_mp@zzzcomputing.com> | 2023-04-22 09:41:49 -0400 |
---|---|---|
committer | Mike Bayer <mike_mp@zzzcomputing.com> | 2023-04-22 12:20:43 -0400 |
commit | d62de424ef738a0cfb9c682d9b25a8644dff5985 (patch) | |
tree | f97f37ade36b3f97cb608757b9b1cbd589fe074b | |
parent | b3e01c30d2610f5454eda5b29f093b65e54230b5 (diff) | |
download | sqlalchemy-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.rst | 9 | ||||
-rw-r--r-- | lib/sqlalchemy/sql/base.py | 22 | ||||
-rw-r--r-- | test/ext/mypy/plain_files/selectables.py | 17 | ||||
-rw-r--r-- | test/sql/test_select.py | 16 |
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"): |