python: expose PanicException (#3455)

pola-rs · May 21, 2022 · be661c7 · be661c7
1 parent e3c7973
commit be661c7
Show file tree

Hide file tree

Showing 5 changed files with 29 additions and 4 deletions.
diff --git a/py-polars/polars/__init__.py b/py-polars/polars/__init__.py
@@ -47,6 +47,7 @@ def version() -> str:
     DuplicateError,
     NoDataError,
     NotFoundError,
+    PanicException,
     SchemaError,
     ShapeError,
 )

diff --git a/py-polars/polars/exceptions.py b/py-polars/polars/exceptions.py
@@ -5,6 +5,7 @@
         DuplicateError,
         NoDataError,
         NotFoundError,
+        PanicException,
         SchemaError,
         ShapeError,
     )
@@ -33,6 +34,9 @@ class ShapeError(Exception):  # type: ignore
     class DuplicateError(Exception):  # type: ignore
         pass
 
+    class PanicException(Exception):  # type: ignore
+        pass
+
 
 __all__ = [
     "ArrowError",
@@ -42,4 +46,5 @@ class DuplicateError(Exception):  # type: ignore
     "SchemaError",
     "ShapeError",
     "DuplicateError",
+    "PanicException",
 ]
diff --git a/py-polars/polars/internals/lazy_functions.py b/py-polars/polars/internals/lazy_functions.py
@@ -1455,30 +1455,38 @@ def select(
 
 @overload
 def struct(
-    exprs: Union[Sequence[Union["pli.Expr", str, "pli.Series"]], "pli.Expr"],
+    exprs: Union[
+        Sequence[Union["pli.Expr", str, "pli.Series"]], "pli.Expr", "pli.Series"
+    ],
     eager: Literal[True],
 ) -> "pli.Series":
     ...
 
 
 @overload
 def struct(
-    exprs: Union[Sequence[Union["pli.Expr", str, "pli.Series"]], "pli.Expr"],
+    exprs: Union[
+        Sequence[Union["pli.Expr", str, "pli.Series"]], "pli.Expr", "pli.Series"
+    ],
     eager: Literal[False],
 ) -> "pli.Expr":
     ...
 
 
 @overload
 def struct(
-    exprs: Union[Sequence[Union["pli.Expr", str, "pli.Series"]], "pli.Expr"],
+    exprs: Union[
+        Sequence[Union["pli.Expr", str, "pli.Series"]], "pli.Expr", "pli.Series"
+    ],
     eager: bool = False,
 ) -> Union["pli.Expr", "pli.Series"]:
     ...
 
 
 def struct(
-    exprs: Union[Sequence[Union["pli.Expr", str, "pli.Series"]], "pli.Expr"],
+    exprs: Union[
+        Sequence[Union["pli.Expr", str, "pli.Series"]], "pli.Expr", "pli.Series"
+    ],
     eager: bool = False,
 ) -> Union["pli.Expr", "pli.Series"]:
     """

diff --git a/py-polars/src/lib.rs b/py-polars/src/lib.rs
@@ -49,6 +49,7 @@ use polars::prelude::Null;
 use polars_core::datatypes::TimeUnit;
 use polars_core::export::arrow::io::ipc::read::read_file_metadata;
 use polars_core::prelude::IntoSeries;
+use pyo3::panic::PanicException;
 use pyo3::types::{PyBool, PyDict, PyFloat, PyInt, PyString};
 #[cfg(target_os = "linux")]
 use tikv_jemallocator::Jemalloc;
@@ -440,6 +441,8 @@ fn polars(py: Python, m: &PyModule) -> PyResult<()> {
         .unwrap();
     m.add("DuplicateError", py.get_type::<DuplicateError>())
         .unwrap();
+    m.add("PanicException", py.get_type::<PanicException>())
+        .unwrap();
     m.add_class::<PySeries>().unwrap();
     m.add_class::<PyDataFrame>().unwrap();
     m.add_class::<PyLazyFrame>().unwrap();

diff --git a/py-polars/tests/test_errors.py b/py-polars/tests/test_errors.py
@@ -49,3 +49,11 @@ def test_not_found_error() -> None:
 def test_string_numeric_comp_err() -> None:
     with pytest.raises(pl.ComputeError, match="cannot compare Utf8 with numeric data"):
         pl.DataFrame({"a": [1.1, 21, 31, 21, 51, 61, 71, 81]}).select(pl.col("a") < "9")
+
+
+def test_panic_exception() -> None:
+    with pytest.raises(
+        pl.PanicException,
+        match=r"""this operation is not implemented/valid for this dtype: .*""",
+    ):
+        pl.struct(pl.Series("a", [1, 2, 3]), eager=True).sort()