From d62de424ef738a0cfb9c682d9b25a8644dff5985 Mon Sep 17 00:00:00 2001 From: Mike Bayer Date: Sat, 22 Apr 2023 09:41:49 -0400 Subject: 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 --- lib/sqlalchemy/sql/base.py | 22 +++++++++++++++++----- 1 file changed, 17 insertions(+), 5 deletions(-) (limited to 'lib/sqlalchemy/sql') 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: -- cgit v1.2.1