fix(drop): support deferred objects in calls to drop

mesejo · Aug 21, 2023 · d27374b · d27374b
1 parent 124f085
commit d27374b
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 0 deletions.
diff --git a/ibis/expr/types/relations.py b/ibis/expr/types/relations.py
@@ -1925,6 +1925,11 @@ def drop(self, *fields: str | Selector) -> Table:
             # no-op if nothing to be dropped
             return self
 
+        fields = tuple(
+            field.resolve(self) if isinstance(field, Deferred) else field
+            for field in fields
+        )
+
         if missing_fields := {f for f in fields if isinstance(f, str)}.difference(
             self.schema().names
         ):

diff --git a/ibis/tests/expr/test_table.py b/ibis/tests/expr/test_table.py
@@ -1626,6 +1626,15 @@ def test_drop():
 
     assert res.equals(t.drop(s.matches("a|b")))
 
+    res = t.drop(_.a)
+    assert res.equals(t.select("b", "c", "d"))
+
+    res = t.drop(_.a, _.b)
+    assert res.equals(t.select("c", "d"))
+
+    res = t.drop(_.a, "b")
+    assert res.equals(t.select("c", "d"))
+
     with pytest.raises(KeyError):
         t.drop("e")