int_range: make datatype check in Rust

pola-rs · Apr 15, 2024 · 6995969 · 6995969
1 parent e392caa
commit 6995969
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 12 deletions.
diff --git a/R/functions__lazy.R b/R/functions__lazy.R
@@ -1313,10 +1313,6 @@ pl_arg_sort_by = function(
 #'   pl$all()
 #' )
 pl_int_range = function(start = 0, end = NULL, step = 1, ..., dtype = pl$Int64) {
-  if (!dtype$is_integer()) {
-    Err_plain("`dtype` must be of type integer") |>
-      unwrap("in pl$int_range():")
-  }
   if (is.null(end)) {
     end = start
     start = 0
@@ -1342,10 +1338,6 @@ pl_int_range = function(start = 0, end = NULL, step = 1, ..., dtype = pl$Int64)
 #'
 #' df$with_columns(int_range = pl$int_ranges("start", "end", dtype = pl$Int16))
 pl_int_ranges = function(start = 0, end = NULL, step = 1, ..., dtype = pl$Int64) {
-  if (!dtype$is_integer()) {
-    Err_plain("`dtype` must be of type integer") |>
-      unwrap("in pl$int_ranges():")
-  }
   if (is.null(end)) {
     end = start
     start = 0

diff --git a/src/rust/src/rlib.rs b/src/rust/src/rlib.rs
@@ -423,7 +423,13 @@ pub fn int_ranges(start: Robj, end: Robj, step: Robj, dtype: Robj) -> RResult<RP
     let start = robj_to!(PLExprCol, start)?;
     let end = robj_to!(PLExprCol, end)?;
     let step = robj_to!(PLExprCol, step)?;
-    let dtype = robj_to!(RPolarsDataType, dtype)?.into();
+    let dtype: pl::DataType = robj_to!(RPolarsDataType, dtype)?.into();
+    if !dtype.is_integer() {
+        return Err(pl::PolarsError::ComputeError(
+            format!("non-integer `dtype` passed to `int_ranges`: {:?}", dtype,).into(),
+        )
+        .into());
+    }
     let mut result = polars::lazy::dsl::int_ranges(start, end, step);
     if dtype != pl::DataType::Int64 {
         result = result.cast(pl::DataType::List(Box::new(dtype)))

diff --git a/tests/testthat/test-lazy_functions.R b/tests/testthat/test-lazy_functions.R
@@ -568,12 +568,12 @@ test_that("pl$int_range() works", {
 
   expect_grepl_error(
     pl$int_range(0, 3, dtype = pl$String) |> as_polars_series(),
-    "must be of type integer"
+    "non-integer `dtype` passed"
   )
 
   expect_grepl_error(
     pl$int_range(0, 3, dtype = pl$Float32) |> as_polars_series(),
-    "must be of type integer"
+    "non-integer `dtype` passed"
   )
 
   # "step" works
@@ -612,7 +612,7 @@ test_that("pl$int_ranges() works", {
 
   expect_grepl_error(
     df$select(int_range = pl$int_ranges("start", "end", dtype = pl$String)),
-    "must be of type integer"
+    "non-integer `dtype` passed"
   )
 
   # "step" works