added slice support for Flags

CHANGELOG.md

@@ -10,6 +10,7 @@ SPDX-License-Identifier: GPL-3.0-or-later
 [List of commits](https://git.ufz.de/rdm-software/saqc/-/compare/v2.3.0...develop)
 ### Added
 - Methods `logicalAnd` and `logicalOr`
+- `Flags` supports slicing and column selection with `list` or a `pd.Index`.
 ### Changed
 ### Removed
 ### Fixed

saqc/core/flags.py

@@ -320,8 +320,28 @@ class Flags:
     # ----------------------------------------------------------------------
     # item access
 
-    def __getitem__(self, key: str) -> pd.Series:
-        return self._data[key].squeeze()
+    def __getitem__(self, key: str | list | pd.Index) -> pd.Series | Flags:
+        if isinstance(key, str):
+            return self._data[key].squeeze()
+
+        if isinstance(key, slice):
+            key = self.columns[key]
+
+        if isinstance(key, (list, pd.Index)):
+            # only copy necessary data
+            data = self._data
+            try:
+                self._data = {}
+                new = self.copy()
+            finally:
+                self._data = data
+            new._data = {k: self._data[k].copy() for k in key}
+            return new
+
+        raise TypeError(
+            "Key must be of type str, list or index of string or slice,"
+            f"not {type(key)}."
+        )
 
     def __setitem__(self, key: SelectT, value: ValueT):
         # force-KW is only internally available

tests/core/test_flags.py

@@ -49,6 +49,7 @@ for d in _arrays:
 
 
 def is_equal(f1, f2):
+    """assert Flags instance equals other"""
     assert f1.columns.equals(f2.columns)
     for c in f1.columns:
         assert test_hist.is_equal(f1.history[c], f2.history[c])
@@ -334,3 +335,50 @@ def test_columns_setter_raises(columns, err):
     )
     with pytest.raises(err):
         flags.columns = columns
+
+
+@pytest.mark.parametrize(
+    "data,key,expected",
+    [
+        (dict(a=[0, 1], b=[]), "a", pd.Series([0, 1], dtype=float)),
+        (dict(a=[0, 1], b=[]), "b", pd.Series([], dtype=float)),
+    ],
+)
+def test__getitem__scalar(data, key, expected):
+    flags = Flags({k: pd.Series(v, dtype=float) for k, v in data.items()})
+    result: pd.Series = flags[key]
+    assert isinstance(result, pd.Series)
+    assert result.equals(expected)
+    # assert copying
+    assert flags[key] is not flags[key]
+
+
+@pytest.mark.parametrize(
+    "data,key,expected",
+    [
+        (dict(a=[0, 1], b=[]), [], dict()),
+        (dict(a=[0, 1], b=[]), ["a"], dict(a=[0, 1])),
+        (dict(a=[0, 1], b=[]), ["a", "b"], dict(a=[0, 1], b=[])),
+        (dict(a=[0, 1], b=[]), pd.Index([]), dict()),
+        (dict(a=[0, 1], b=[]), pd.Index(["a"]), dict(a=[0, 1])),
+        (dict(a=[0, 1], b=[]), pd.Index(["a", "b"]), dict(a=[0, 1], b=[])),
+        (dict(a=[0, 1], b=[]), slice(None), dict(a=[0, 1], b=[])),
+        (dict(a=[0, 1], b=[]), slice(0, 1), dict(a=[0, 1])),
+        (dict(a=[0, 1], b=[]), slice(1, 99), dict(b=[])),
+        (dict(a=[0, 1], b=[]), slice(5, 99), dict()),
+    ],
+)
+def test__getitem__listlike_and_slice(data, key, expected):
+    flags = Flags({k: pd.Series(v, dtype=float) for k, v in data.items()})
+    result: Flags = flags[key]
+    assert isinstance(result, Flags)
+    # assert that a new Flags object was created
+    assert flags[key] is not flags[key]
+    # assert that internal data is copied
+    if len(result):
+        left = result._data[result.columns[0]]
+        right = flags._data[result.columns[0]]
+        assert left is not right
+
+    expected = Flags({k: pd.Series(v, dtype=float) for k, v in expected.items()})
+    is_equal(result, expected)