More narwhals

innobi · Aug 17, 2024 · e76be41 · e76be41
1 parent f00413e
commit e76be41
Showing 1 changed file with 2 additions and 8 deletions.
diff --git a/tests/conftest.py b/tests/conftest.py
@@ -484,15 +484,9 @@ def drop_columns(frame, columns):
         return frame.drop(columns)
 
     @staticmethod
+    @nw.narwhalify
     def select_columns(frame, columns):
-        if isinstance(frame, pd.DataFrame):
-            return frame[columns]
-        elif isinstance(frame, pa.Table):
-            return frame.select(columns)
-        elif isinstance(frame, pl.DataFrame):
-            return frame.select(columns)
-        else:
-            raise NotImplementedError("select_columns not implemented for type")
+        frame.select(columns)
 
     @staticmethod
     @nw.narwhalify