From 536eef793644ea678b46d8b5bde11bd491ae25c7 Mon Sep 17 00:00:00 2001
From: Bert Palm <bert.palm@ufz.de>
Date: Mon, 27 Feb 2023 12:43:28 +0100
Subject: [PATCH] added slice support for Flags, address #355
---
saqc/core/flags.py | 24 ++++++++++++++++++--
tests/core/test_flags.py | 48 ++++++++++++++++++++++++++++++++++++++++
2 files changed, 70 insertions(+), 2 deletions(-)
diff --git a/saqc/core/flags.py b/saqc/core/flags.py
index 3e7e7dcc0..b1ec0b580 100644
--- a/saqc/core/flags.py
+++ b/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
diff --git a/tests/core/test_flags.py b/tests/core/test_flags.py
index 74e3f4d90..6b80d4ddb 100644
--- a/tests/core/test_flags.py
+++ b/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)
--
GitLab