Repository: numpy/numpy
Pull request: BUG: Fix strange behavior of infinite-step-size/underflow-case in arange
This closes https://github.com/numpy/numpy/issues/10206.

The problem is
```
>>> np.arange(0, 1, 1e10)
array([0.])
>>> np.arange(0, 1, 1e100)
array([0.])
>>> np.arange(0, 1, 1e1000)  # overflows
array([], dtype=float64)
>>> np.arange(0, 1, np.inf)
array([], dtype=float64)
```
Where the step size becomes infinite.

Problem statement:
BUG: Strange behavior of infinite step size in arange
```python
>>> np.arange(0, 1, 1e10)
array([0.])
>>> np.arange(0, 1, 1e100)
array([0.])
>>> np.arange(0, 1, 1e1000)  # overflows
array([], dtype=float64)
>>> np.arange(0, 1, np.inf)
array([], dtype=float64)
```

I think the last two cases should match the others.

Hints from issue discussion:

Commit message: Add zero PyObject pointer null-check

Patch to review:
diff --git a/numpy/core/src/multiarray/ctors.c b/numpy/core/src/multiarray/ctors.c
index f4236f36dfa5..75b99931f369 100644
--- a/numpy/core/src/multiarray/ctors.c
+++ b/numpy/core/src/multiarray/ctors.c
@@ -2968,11 +2968,26 @@ PyArray_Arange(double start, double stop, double step, int type_num)
     PyArray_ArrFuncs *funcs;
     PyObject *obj;
     int ret;
+    double delta, tmp_len;
     NPY_BEGIN_THREADS_DEF;
 
-    length = _arange_safe_ceil_to_intp((stop - start)/step);
-    if (error_converting(length)) {
-        return NULL;
+    delta = stop - start;
+    tmp_len = delta/step;
+
+    /* Underflow and divide-by-inf check */
+    if (tmp_len == 0.0 && delta != 0.0) {
+        if (npy_signbit(tmp_len)) {
+            length = -1.0;
+        }
+        else {
+            length = 1.0;
+        }
+    }
+    else {
+        length = _arange_safe_ceil_to_intp(tmp_len);
+        if (error_converting(length)) {
+            return NULL;
+        }
     }
 
     if (length <= 0) {
@@ -3036,11 +3051,19 @@ static npy_intp
 _calc_length(PyObject *start, PyObject *stop, PyObject *step, PyObject **next, int cmplx)
 {
     npy_intp len, tmp;
-    PyObject *val;
+    PyObject *zero, *val;
+    int next_is_nonzero, val_is_zero;
     double value;
 
+    zero = PyInt_FromLong(0);
+    if (!zero) {
+        return -1;
+    }
+
     *next = PyNumber_Subtract(stop, start);
     if (!(*next)) {
+        Py_DECREF(zero);
+
         if (PyTuple_Check(stop)) {
             PyErr_Clear();
             PyErr_SetString(PyExc_TypeError,
@@ -3049,12 +3072,20 @@ _calc_length(PyObject *start, PyObject *stop, PyObject *step, PyObject **next, i
         }
         return -1;
     }
+
+    next_is_nonzero = PyObject_RichCompareBool(*next, zero, Py_NE);
     val = PyNumber_TrueDivide(*next, step);
     Py_DECREF(*next);
     *next = NULL;
+
     if (!val) {
+        Py_DECREF(zero);
         return -1;
     }
+
+    val_is_zero = PyObject_RichCompareBool(val, zero, Py_EQ);
+    Py_DECREF(zero);
+
     if (cmplx && PyComplex_Check(val)) {
         value = PyComplex_RealAsDouble(val);
         if (error_converting(value)) {
@@ -3083,11 +3114,24 @@ _calc_length(PyObject *start, PyObject *stop, PyObject *step, PyObject **next, i
         if (error_converting(value)) {
             return -1;
         }
-        len = _arange_safe_ceil_to_intp(value);
-        if (error_converting(len)) {
-            return -1;
+
+        /* Underflow and divide-by-inf check */
+        if (val_is_zero && next_is_nonzero) {
+            if (npy_signbit(value)) {
+                len = -1.0;
+            }
+            else {
+                len = 1.0;
+            }
+        }
+        else {
+            len = _arange_safe_ceil_to_intp(value);
+            if (error_converting(len)) {
+                return -1;
+            }
         }
     }
+
     if (len > 0) {
         *next = PyNumber_Add(start, step);
         if (!*next) {
diff --git a/numpy/core/tests/test_regression.py b/numpy/core/tests/test_regression.py
index a3b011454d48..ee18060d6618 100644
--- a/numpy/core/tests/test_regression.py
+++ b/numpy/core/tests/test_regression.py
@@ -224,6 +224,23 @@ def test_arange_endian(self):
         x = np.arange(10, dtype='>f8')
         assert_array_equal(ref, x)
 
+    def test_arange_inf_step(self):
+        ref = np.arange(0, 1, 10)
+        x = np.arange(0, 1, np.inf)
+        assert_array_equal(ref, x)
+
+        ref = np.arange(0, 1, -10)
+        x = np.arange(0, 1, -np.inf)
+        assert_array_equal(ref, x)
+
+        ref = np.arange(0, -1, -10)
+        x = np.arange(0, -1, -np.inf)
+        assert_array_equal(ref, x)
+
+        ref = np.arange(0, -1, 10)
+        x = np.arange(0, -1, np.inf)
+        assert_array_equal(ref, x)
+
     def test_argmax(self):
         # Ticket #119
         a = np.random.normal(0, 1, (4, 5, 6, 7, 8))


Code review comments:
diff_hunk: @@ -3036,11 +3051,19 @@ static npy_intp
 _calc_length(PyObject *start, PyObject *stop, PyObject *step, PyObject **next, int cmplx)
 {
     npy_intp len, tmp;
-    PyObject *val;
+    PyObject *zero, *val;
+    int next_is_nonzero, val_is_zero;
     double value;
 
+    zero = PyInt_FromLong(0);
+    if (!zero) {
+        return -1;
+    }
+
     *next = PyNumber_Subtract(stop, start);
     if (!(*next)) {
+        Py_DECREF(zero);
path: numpy/core/src/multiarray/ctors.c
text: @user1:
If you move the allocation of `zero` down below this `if`, then you don't need to clear it up here

@author:
@user2-wieser If I move down `zero` allocation, we need to `DECREF` `*next` and set `NULL` to `*next` when the `zero` allocation failed. IMO, this is complicated.

@user1:
Fair


@user1:
>  and set NULL to *next when the zero allocation failed. 

You already need to do this

diff_hunk: @@ -3036,11 +3051,19 @@ static npy_intp
 _calc_length(PyObject *start, PyObject *stop, PyObject *step, PyObject **next, int cmplx)
 {
     npy_intp len, tmp;
-    PyObject *val;
+    PyObject *zero, *val;
+    int next_is_nonzero, val_is_zero;
     double value;
 
+    zero = PyInt_FromLong(0);
+    if (!zero) {
+        return -1;
path: numpy/core/src/multiarray/ctors.c
text: @user1:
I think the problem is that `*next` is not yet set here, which causes the `Py_XDECREF(next);` in the calling function to fail.

@user1:

Also, that `Py_DECREF(next);` should be `Py_XDECREF`...