diff --git a/core/eolearn/core/eodata.py b/core/eolearn/core/eodata.py index a1ce4d697..19dd64ea5 100644 --- a/core/eolearn/core/eodata.py +++ b/core/eolearn/core/eodata.py @@ -352,19 +352,22 @@ def __setitem__( return self.__setattr__(FeatureType(feature_type).value, value, feature_name=feature_name) - def __delitem__(self, feature: Tuple[FeatureType, str]) -> None: + def __delitem__(self, feature: FeatureSpec) -> None: """Deletes the selected feature. :param feature: EOPatch feature """ self._check_tuple_key(feature) feature_type, feature_name = feature - del self[feature_type][feature_name] + if feature_type in [FeatureType.BBOX, FeatureType.TIMESTAMP]: + self.reset_feature_type(feature_type) + else: + del self[feature_type][feature_name] @staticmethod def _check_tuple_key(key: tuple) -> None: """A helper function that checks a tuple, which should hold (feature_type, feature_name).""" - if len(key) != 2: + if not isinstance(key, (tuple, list)) or len(key) != 2: raise ValueError(f"Given element should be a tuple of (feature_type, feature_name), but {key} found.") def __eq__(self, other: object) -> bool: diff --git a/core/eolearn/tests/test_eodata.py b/core/eolearn/tests/test_eodata.py index fd95f37ea..042efa902 100644 --- a/core/eolearn/tests/test_eodata.py +++ b/core/eolearn/tests/test_eodata.py @@ -163,9 +163,10 @@ def test_simplified_feature_operations() -> None: (FeatureType.MASK, "ones"), (FeatureType.MASK_TIMELESS, "threes"), (FeatureType.META_INFO, "beep"), + (FeatureType.BBOX, None), ], ) -def test_delete_existing_feature(feature: Tuple[FeatureType, str], mini_eopatch: EOPatch) -> None: +def test_delete_existing_feature(feature: FeatureSpec, mini_eopatch: EOPatch) -> None: old = mini_eopatch.copy(deep=True) del mini_eopatch[feature] @@ -173,6 +174,7 @@ def test_delete_existing_feature(feature: Tuple[FeatureType, str], mini_eopatch: for old_feature in old.get_features(): if old_feature != feature: + # this also works for BBox :D assert_array_equal(old[old_feature], mini_eopatch[old_feature])