take advantage of try_get_column_index

pola-rs · Jul 26, 2024 · cc0440b · cc0440b
1 parent 1484ef8
commit cc0440b
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 8 deletions.
diff --git a/py-polars/polars/dataframe/frame.py b/py-polars/polars/dataframe/frame.py
@@ -4755,12 +4755,9 @@ def get_column_index(self, name: str) -> int:
         >>> df.get_column_index("ham")
         2
         >>> df.get_column_index("sandwich")  # doctest: +SKIP
-        ColumnNotFoundError: 'sandwich' not found in DataFrame
+        ColumnNotFoundError: sandwich
         """
-        if (idx := self._df.get_column_index(name)) is None:
-            msg = f"{name!r} not found in DataFrame"
-            raise ColumnNotFoundError(msg)
-        return idx
+        return self._df.get_column_index(name)
 
     def replace_column(self, index: int, column: Series) -> DataFrame:
         """

diff --git a/py-polars/src/dataframe/general.rs b/py-polars/src/dataframe/general.rs
@@ -226,8 +226,11 @@ impl PyDataFrame {
         }
     }
 
-    pub fn get_column_index(&self, name: &str) -> Option<usize> {
-        self.df.get_column_index(name)
+    pub fn get_column_index(&self, name: &str) -> PyResult<usize> {
+        Ok(self
+            .df
+            .try_get_column_index(name)
+            .map_err(PyPolarsErr::from)?)
     }
 
     pub fn get_column(&self, name: &str) -> PyResult<PySeries> {

diff --git a/py-polars/tests/unit/dataframe/test_df.py b/py-polars/tests/unit/dataframe/test_df.py
@@ -2848,5 +2848,5 @@ def test_get_column_index() -> None:
     assert df.get_column_index("actual") == 0
     assert df.get_column_index("expected") == 1
 
-    with pytest.raises(ColumnNotFoundError, match="'missing' not found in DataFrame"):
+    with pytest.raises(ColumnNotFoundError, match="missing"):
         df.get_column_index("missing")