diff --git a/dvc/compare.py b/dvc/compare.py index 7808359b2..9384d5f1f 100644 --- a/dvc/compare.py +++ b/dvc/compare.py @@ -12,6 +12,7 @@ from typing import ( Mapping, MutableSequence, Sequence, + Set, Tuple, Union, overload, @@ -181,6 +182,35 @@ class TabularData(MutableSequence[Sequence["CellT"]]): {k: self._columns[k][i] for k in keys} for i in range(len(self)) ] + def dropna(self, axis: str = "rows"): + if axis not in ["rows", "cols"]: + raise ValueError( + f"Invalid 'axis' value {axis}." + "Choose one of ['rows', 'cols']" + ) + + to_drop: Set = set() + for n_row, row in enumerate(self): + for n_col, col in enumerate(row): + if col == self._fill_value: + if axis == "rows": + to_drop.add(n_row) + break + else: + to_drop.add(self.keys()[n_col]) + + if axis == "rows": + for name in self.keys(): + self._columns[name] = Column( + [ + x + for n, x in enumerate(self._columns[name]) + if n not in to_drop + ] + ) + else: + self.drop(*to_drop) + def _normalize_float(val: float, precision: int): return f"{val:.{precision}g}" diff --git a/tests/unit/test_tabular_data.py b/tests/unit/test_tabular_data.py index 0b815eae7..9631c4362 100644 --- a/tests/unit/test_tabular_data.py +++ b/tests/unit/test_tabular_data.py @@ -177,3 +177,36 @@ def test_row_from_dict(): ["", "", "value3", "value4", "", ""], ["", "", "value3", "", "value5", "value6"], ] + + +@pytest.mark.parametrize( + "axis,expected", + [ + ( + "rows", + [ + ["foo", "bar", "foobar"], + ], + ), + ("cols", [["foo"], ["foo"], ["foo"]]), + ], +) +def test_dropna(axis, expected): + td = TabularData(["col-1", "col-2", "col-3"]) + td.extend([["foo"], ["foo", "bar"], ["foo", "bar", "foobar"]]) + assert list(td) == [ + ["foo", "", ""], + ["foo", "bar", ""], + ["foo", "bar", "foobar"], + ] + + td.dropna(axis) + + assert list(td) == expected + + +def test_dropna_invalid_axis(): + td = TabularData(["col-1", "col-2", "col-3"]) + + with pytest.raises(ValueError, match="Invalid 'axis' value foo."): + td.dropna("foo")