Fix arange crashing when start, stop or step is nan (#605)

v923z · Apr 27, 2023 · 8c3e105 · 8c3e105
1 parent ad1a1c5
commit 8c3e105
Show file tree

Hide file tree

Showing 4 changed files with 12 additions and 3 deletions.
diff --git a/code/numpy/create.c b/code/numpy/create.c
@@ -158,6 +158,10 @@ mp_obj_t create_arange(size_t n_args, const mp_obj_t *pos_args, mp_map_t *kw_arg
         mp_raise_msg(&mp_type_ZeroDivisionError, MP_ERROR_TEXT("divide by zero"));
     }
 
+    if(isnan(start) || isnan(stop) || isnan(step)) {
+        mp_raise_ValueError(translate("arange: cannot compute length"));
+    }
+
     ndarray_obj_t *ndarray;
     if((stop - start)/step <= 0) {
         ndarray = ndarray_new_linear_array(0, dtype);
@@ -254,8 +258,6 @@ mp_obj_t create_concatenate(size_t n_args, const mp_obj_t *pos_args, mp_map_t *k
     size_t *shape = m_new0(size_t, ULAB_MAX_DIMS);
     mp_obj_tuple_t *ndarrays = MP_OBJ_TO_PTR(args[0].u_obj);
 
-    // first check, whether
-
     for(uint8_t i = 0; i < ndarrays->len; i++) {
         if(!mp_obj_is_type(ndarrays->items[i], &ulab_ndarray_type)) {
             mp_raise_ValueError(translate("only ndarrays can be concatenated"));

diff --git a/code/ulab.c b/code/ulab.c
@@ -33,7 +33,7 @@
 #include "user/user.h"
 #include "utils/utils.h"
 
-#define ULAB_VERSION 6.0.8
+#define ULAB_VERSION 6.0.9
 #define xstr(s) str(s)
 #define str(s) #s
 

diff --git a/tests/2d/numpy/arange.py b/tests/2d/numpy/arange.py
@@ -17,3 +17,9 @@
     np.arange(0, 10, 0)
 except ZeroDivisionError as e:
     print('ZeroDivisionError: ', e)
+
+# test for NAN length exception
+try:
+    np.arange(0, np.nan)
+except ValueError as e:
+    print('ValueError: ', e)
diff --git a/tests/2d/numpy/arange.py.exp b/tests/2d/numpy/arange.py.exp
@@ -19,3 +19,4 @@ array([2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0], dtype=float64)
 array([2.0, 5.0, 8.0], dtype=float64)
 array([], dtype=float64)
 ZeroDivisionError:  divide by zero
+ValueError:  arange: cannot compute length