From d4861d535ab868e59a5b602846ef70ebe6526b75 Mon Sep 17 00:00:00 2001
From: David de la Iglesia Castro <daviddelaiglesiacastro@gmail.com>
Date: Thu, 20 Jan 2022 20:02:53 +0100
Subject: [PATCH] TabularData: Simplify drop internals.

Update `subset` type hint.
---
 dvc/commands/experiments/show.py       |  2 +-
 dvc/compare.py                         | 14 +++++++-------
 tests/unit/command/test_experiments.py |  4 ++--
 3 files changed, 10 insertions(+), 10 deletions(-)

diff --git a/dvc/commands/experiments/show.py b/dvc/commands/experiments/show.py
index ead3ed542..87a0469dc 100644
--- a/dvc/commands/experiments/show.py
+++ b/dvc/commands/experiments/show.py
@@ -391,7 +391,7 @@ def show_experiments(
     td.drop(*cols_to_drop)
 
     if pcp:
-        subset = [x for x in td.keys() if x != "Experiment"]
+        subset = {x for x in td.keys() if x != "Experiment"}
         td.dropna(
             "rows",
             how="all",
diff --git a/dvc/compare.py b/dvc/compare.py
index 9037cd637..aa3a2cd45 100644
--- a/dvc/compare.py
+++ b/dvc/compare.py
@@ -137,13 +137,10 @@ class TabularData(MutableSequence[Sequence["CellT"]]):
         return len(self.columns), len(self)
 
     def drop(self, *col_names: str) -> None:
-        to_remove = set()
         for col in col_names:
             if not self.is_protected(col):
-                to_remove.add(col)
-        for col in to_remove:
-            self._keys.remove(col)
-            self._columns.pop(col)
+                self._keys.remove(col)
+                self._columns.pop(col)
 
     def rename(self, from_col_name: str, to_col_name: str) -> None:
         self._columns[to_col_name] = self._columns.pop(from_col_name)
@@ -215,7 +212,10 @@ class TabularData(MutableSequence[Sequence["CellT"]]):
         ]
 
     def dropna(
-        self, axis: str = "rows", how="any", subset: Optional[List] = None
+        self,
+        axis: str = "rows",
+        how="any",
+        subset: Optional[Iterable[str]] = None,
     ):
         if axis not in ["rows", "cols"]:
             raise ValueError(
@@ -266,7 +266,7 @@ class TabularData(MutableSequence[Sequence["CellT"]]):
     def drop_duplicates(
         self,
         axis: str = "rows",
-        subset: Optional[List] = None,
+        subset: Optional[Iterable[str]] = None,
         ignore_empty: bool = True,
     ):
         if axis not in ["rows", "cols"]:
diff --git a/tests/unit/command/test_experiments.py b/tests/unit/command/test_experiments.py
index df3bb3feb..966070cfb 100644
--- a/tests/unit/command/test_experiments.py
+++ b/tests/unit/command/test_experiments.py
@@ -791,8 +791,8 @@ def test_show_experiments_pcp(tmp_dir, mocker):
     show_experiments(all_experiments, pcp=True)
 
     td.drop.assert_called_with("Created")
-    td.dropna.assert_called_with("rows", how="all", subset=[])
-    td.drop_duplicates.assert_called_with("rows", subset=[])
+    td.dropna.assert_called_with("rows", how="all", subset=set())
+    td.drop_duplicates.assert_called_with("rows", subset=set())
 
     kwargs = td.to_parallel_coordinates.call_args[1]