diff --git a/changelog.d/contained-labels-pandas3.rst b/changelog.d/contained-labels-pandas3.rst new file mode 100644 index 00000000..f717d39e --- /dev/null +++ b/changelog.d/contained-labels-pandas3.rst @@ -0,0 +1,2 @@ +**Fix get_contained_labels with pandas 3** + ``Annotation.get_contained_labels`` no longer raises a ``TypeError`` with pandas 3 for annotations without ``label_store`` values, such as those read with ``rdann`` by default. diff --git a/src/wfdb/io/annotation.py b/src/wfdb/io/annotation.py index 6c3e3b0f..5914785c 100644 --- a/src/wfdb/io/annotation.py +++ b/src/wfdb/io/annotation.py @@ -1350,12 +1350,12 @@ def get_contained_labels(self, inplace=True): counts = np.unique(self.label_store, return_counts=True) elif self.symbol is not None: index_vals = set(self.symbol) - label_map = label_map.set_index(label_map["symbol"].values) + label_map = label_map.set_index("symbol", drop=False) reset_index = True counts = np.unique(self.symbol, return_counts=True) elif self.description is not None: index_vals = set(self.description) - label_map = label_map.set_index(label_map["description"].values) + label_map = label_map.set_index("description", drop=False) reset_index = True counts = np.unique(self.description, return_counts=True) else: diff --git a/tests/test_annotation.py b/tests/test_annotation.py index a0e901fd..37a9a8b6 100644 --- a/tests/test_annotation.py +++ b/tests/test_annotation.py @@ -342,3 +342,18 @@ def test_aux_note_none(temp_path): open(os.path.join(temp_path, "empty.atr"), "rb") as f2, ): assert f1.read() == f2.read() + + +def test_contained_labels_from_symbols(): + """Contained labels can be found from symbols alone""" + ann = wfdb.rdann("sample-data/100", "atr") + assert ann.label_store is None + labels = ann.get_contained_labels(inplace=False) + assert list(labels.columns) == [ + "label_store", + "symbol", + "description", + "n_occurrences", + ] + assert labels.loc[1, "symbol"] == "N" + assert labels["n_occurrences"].sum() == len(ann.sample)