From 8d42733e1b5b0b49cb39c9f3717ac6eede31a220 Mon Sep 17 00:00:00 2001 From: Tim Sawicki Date: Tue, 16 Jul 2024 15:43:10 -0400 Subject: [PATCH] ignore empty lists when dropping --- quixstreams/dataframe/dataframe.py | 2 ++ .../test_quixstreams/test_dataframe/test_dataframe.py | 10 ++++++++++ 2 files changed, 12 insertions(+) diff --git a/quixstreams/dataframe/dataframe.py b/quixstreams/dataframe/dataframe.py index 8840bfe3e..9dc0bf034 100644 --- a/quixstreams/dataframe/dataframe.py +++ b/quixstreams/dataframe/dataframe.py @@ -933,6 +933,8 @@ def drop(self, columns: Union[str, List[str]]) -> Self: :return: a new StreamingDataFrame instance """ if isinstance(columns, list): + if not columns: + return self if not all(isinstance(s, str) for s in columns): raise TypeError(f"column list must contain strings only") elif isinstance(columns, str): diff --git a/tests/test_quixstreams/test_dataframe/test_dataframe.py b/tests/test_quixstreams/test_dataframe/test_dataframe.py index ec43fd37b..20b52a566 100644 --- a/tests/test_quixstreams/test_dataframe/test_dataframe.py +++ b/tests/test_quixstreams/test_dataframe/test_dataframe.py @@ -392,6 +392,16 @@ def test_drop_invalid_columns(self, dataframe_factory, columns): with pytest.raises(TypeError): sdf.drop(columns) + def test_drop_empty_list(self, dataframe_factory): + """ + Dropping an empty list is ignored entirely. + """ + sdf = dataframe_factory() + pre_drop_stream = sdf.stream.tree() + sdf = sdf.drop([]) + post_drop_stream = sdf.stream.tree() + assert pre_drop_stream == post_drop_stream + class TestStreamingDataFrameApplyExpand: def test_apply_expand(self, dataframe_factory):