From 9fc5247937e6bde9ca7e798807c9d864c3963c58 Mon Sep 17 00:00:00 2001
From: Ritchie Vink <ritchie46@gmail.com>
Date: Tue, 2 Aug 2022 15:37:53 +0200
Subject: [PATCH] python: improve error (#4223)

---
 py-polars/polars/internals/frame.py  | 11 +++++++++--
 py-polars/polars/internals/series.py |  5 ++++-
 py-polars/src/lazy/apply.rs          | 10 ++++------
 py-polars/src/lazy/dsl.rs            |  3 ++-
 py-polars/tests/test_df.py           |  4 ++--
 py-polars/tests/test_errors.py       | 23 +++++++++++++++++++++++
 6 files changed, 44 insertions(+), 12 deletions(-)

diff --git a/py-polars/polars/internals/frame.py b/py-polars/polars/internals/frame.py
index 2acbdce49082..341f5131b114 100644
--- a/py-polars/polars/internals/frame.py
+++ b/py-polars/polars/internals/frame.py
@@ -1993,7 +1993,10 @@ def __getitem__(
                 )
 
         # if no data has been returned, the operation is not supported
-        raise NotImplementedError
+        raise ValueError(
+            f"Cannot __getitem__ on DataFrame with item: '{item}'"
+            f" of type: '{type(item)}'."
+        )
 
     def __setitem__(
         self, key: str | list | tuple[Any, str | int], value: Any
@@ -2048,7 +2051,11 @@ def __setitem__(
             elif isinstance(col_selection, str):
                 self.replace(col_selection, s)
         else:
-            raise NotImplementedError
+            raise ValueError(
+                f"Cannot __setitem__ on DataFrame with key: '{key}' "
+                f"of type: '{type(key)}' and value: '{value}' "
+                f"of type: '{type(value)}'."
+            )
 
     def __len__(self) -> int:
         return self.height
diff --git a/py-polars/polars/internals/series.py b/py-polars/polars/internals/series.py
index 863a09f2488e..2e507e8e2b13 100644
--- a/py-polars/polars/internals/series.py
+++ b/py-polars/polars/internals/series.py
@@ -579,7 +579,10 @@ def __getitem__(self, item: int | Series | range | slice | np.ndarray) -> Any:
         if isinstance(item, slice):
             return PolarsSlice(self).apply(item)
 
-        raise NotImplementedError
+        raise ValueError(
+            f"Cannot __getitem__ on Series of dtype: '{self.dtype}' "
+            f"with argument: '{item}' of type: '{type(item)}'."
+        )
 
     def __setitem__(
         self, key: int | Series | np.ndarray | list | tuple, value: Any
diff --git a/py-polars/src/lazy/apply.rs b/py-polars/src/lazy/apply.rs
index fd54b7af19a4..e8ed3dedd340 100644
--- a/py-polars/src/lazy/apply.rs
+++ b/py-polars/src/lazy/apply.rs
@@ -49,7 +49,7 @@ pub(crate) fn call_lambda_with_series(
     s: Series,
     lambda: &PyObject,
     polars_module: &PyObject,
-) -> PyObject {
+) -> PyResult<PyObject> {
     let pypolars = polars_module.cast_as::<PyModule>(py).unwrap();
 
     // create a PySeries struct/object for Python
@@ -61,10 +61,7 @@ pub(crate) fn call_lambda_with_series(
         .call1((pyseries,))
         .unwrap();
     // call the lambda and get a python side Series wrapper
-    match lambda.call1(py, (python_series_wrapper,)) {
-        Ok(pyobj) => pyobj,
-        Err(e) => panic!("python apply failed: {}", e.value(py)),
-    }
+    lambda.call1(py, (python_series_wrapper,))
 }
 
 /// A python lambda taking two Series
@@ -159,7 +156,8 @@ pub fn map_single(
         let py = gil.python();
 
         // this is a python Series
-        let out = call_lambda_with_series(py, s.clone(), &lambda, &pypolars);
+        let out = call_lambda_with_series(py, s.clone(), &lambda, &pypolars)
+            .map_err(|e| PolarsError::ComputeError(format!("{e}").into()))?;
 
         Ok(out.to_series(py, &pypolars, s.name()))
     };
diff --git a/py-polars/src/lazy/dsl.rs b/py-polars/src/lazy/dsl.rs
index cbfbbfbe4530..a5087baab052 100644
--- a/py-polars/src/lazy/dsl.rs
+++ b/py-polars/src/lazy/dsl.rs
@@ -912,7 +912,8 @@ impl PyExpr {
             let gil = Python::acquire_gil();
             let py = gil.python();
 
-            let out = call_lambda_with_series(py, s.clone(), &lambda, &pypolars);
+            let out = call_lambda_with_series(py, s.clone(), &lambda, &pypolars)
+                .expect("python function failed");
             match out.getattr(py, "_s") {
                 Ok(pyseries) => {
                     let pyseries = pyseries.extract::<PySeries>(py).unwrap();
diff --git a/py-polars/tests/test_df.py b/py-polars/tests/test_df.py
index 532389815127..f09281782093 100644
--- a/py-polars/tests/test_df.py
+++ b/py-polars/tests/test_df.py
@@ -612,7 +612,7 @@ def test_set() -> None:
             df[(1, 2, 3)] = 1  # type: ignore[index]
 
         # we cannot index with any type, such as bool
-        with pytest.raises(NotImplementedError):
+        with pytest.raises(ValueError):
             df[True] = 1  # type: ignore[index]
 
 
@@ -1660,7 +1660,7 @@ def test_get_item() -> None:
 
     # note that we cannot use floats (even if they could be casted to integer without
     # loss)
-    with pytest.raises(NotImplementedError):
+    with pytest.raises(ValueError):
         _ = df[np.array([1.0])]
 
     # using boolean masks with numpy is deprecated
diff --git a/py-polars/tests/test_errors.py b/py-polars/tests/test_errors.py
index c9d79dc8189e..5095b200c98e 100644
--- a/py-polars/tests/test_errors.py
+++ b/py-polars/tests/test_errors.py
@@ -117,3 +117,26 @@ def test_not_found_on_rename() -> None:
                 "does_not_exist": "exists",
             }
         )
+
+
+@typing.no_type_check
+def test_getitem_errs() -> None:
+    df = pl.DataFrame({"a": [1, 2, 3]})
+
+    with pytest.raises(
+        ValueError,
+        match=r"Cannot __getitem__ on DataFrame with item: "
+        r"'{'some'}' of type: '<class 'set'>'.",
+    ):
+        df[{"some"}]
+
+    with pytest.raises(
+        ValueError,
+        match=r"Cannot __getitem__ on Series of dtype: "
+        r"'<class 'polars.datatypes.Int64'>' with argument: "
+        r"'{'strange'}' of type: '<class 'set'>'.",
+    ):
+        df["a"][{"strange"}]
+
+    with pytest.raises(ValueError, match="Cannot __setitem__ on DataFrame with key:.*"):
+        df[{"some"}] = "foo"