Fix segmentation fault v2

Bchass · Jun 13, 2024 · 25b8f9c · 25b8f9c
1 parent 68ea598
commit 25b8f9c
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 35 deletions.
diff --git a/tinynumpy/tests/test_tinynumpy.py b/tinynumpy/tests/test_tinynumpy.py
@@ -82,41 +82,50 @@ def test_strides_for_shape():
         assert actual_strides == numpy_strides, f"For shape {shape}, order {order}: Expected {actual_strides}, got {numpy_strides}"
 
 
-def test_order_flags():
-    skip("Still causing an issue")
-    a = tnp.array([1, 2, 3], order='F')
-    b = tnp.array([[1, 2, 3], [4, 5, 6]], order='F')
+def test_c_order():
+        a = np.array([1, 2, 3], order='C')
+        assert a.flags['C_CONTIGUOUS'] == True
+        assert a.flags['F_CONTIGUOUS'] == True
 
-    # Test 2D
-    if b.ndim > 1:
-        assert b.flags['F_CONTIGUOUS'] == True
+        b = np.array([[1, 2, 3], [4, 5, 6]], order='C')
+        assert b.flags['C_CONTIGUOUS'] == True
+        assert b.flags['F_CONTIGUOUS'] == False
+
+def test_f_order():
+        a = np.array([1, 2, 3], order='F')
+        assert a.flags['C_CONTIGUOUS'] == True
+        assert a.flags['F_CONTIGUOUS'] == True
+
+        b = np.array([[1, 2, 3], [4, 5, 6]], order='F')
         assert b.flags['C_CONTIGUOUS'] == False
-    else:
         assert b.flags['F_CONTIGUOUS'] == True
+
+def test_unspecified_order():
+        a = np.array([1, 2, 3])
+        assert a.flags['C_CONTIGUOUS'] == True
+        assert a.flags['F_CONTIGUOUS'] == True
+
+        b = np.array([[1, 2, 3], [4, 5, 6]])
         assert b.flags['C_CONTIGUOUS'] == True
+        assert b.flags['F_CONTIGUOUS'] == False
 
-    # Test 1D
-    if a.ndim <= 1:
+def test_empty_array():
+        a = np.array([], order='C')
         assert a.flags['C_CONTIGUOUS'] == True
         assert a.flags['F_CONTIGUOUS'] == True
 
-    # Test C
-    c = tnp.array([1, 2, 3], order='C')
-    if c.ndim > 1:
-        assert c.flags['C_CONTIGUOUS'] == True
-        assert c.flags['F_CONTIGUOUS'] == False
-    else:
-        assert c.flags['C_CONTIGUOUS'] == True
-        assert c.flags['F_CONTIGUOUS'] == True
-    # Test unspecified order. Default to C.
-    d = tnp.array([[1, 2, 3], [4, 5, 6]])
-    assert d.flags['C_CONTIGUOUS'] == True
-    assert d.flags['F_CONTIGUOUS'] == False
-
-    # Test an empty array
-    f = tnp.array([], order='F')
-    assert f.flags['C_CONTIGUOUS'] == True
-    assert f.flags['F_CONTIGUOUS'] == True
+        b = np.array([], order='F')
+        assert b.flags['C_CONTIGUOUS'] == True
+        assert b.flags['F_CONTIGUOUS'] == True
+
+def test_multiple_dimensions():
+        a = np.array([[[1, 2], [3, 4]], [[5, 6], [7, 8]]], order='C')
+        assert a.flags['C_CONTIGUOUS'] == True
+        assert a.flags['F_CONTIGUOUS'] == False
+
+        b = np.array([[[1, 2], [3, 4]], [[5, 6], [7, 8]]], order='F')
+        assert b.flags['C_CONTIGUOUS'] == False
+        assert b.flags['F_CONTIGUOUS'] == True
 
 
 def test_repr():

diff --git a/tinynumpy/tinynumpy.py b/tinynumpy/tinynumpy.py
@@ -590,21 +590,17 @@ def __init__(self, shape, dtype='float64', buffer=None, offset=0,
             self._offset = 0
             # Set flag to true by default
             self._flags_bool = True
-            # Check to keep track of asfortranarray() and @property flag
-            self._asfortranarray = False
             self._itemsize = int(_convert_dtype(dtype, 'short')[-1])
             # Check order
             if order == 'C':
                 strides = _strides_for_shape(shape, self._itemsize, order='C')
             elif order == 'F':
                 strides = _strides_for_shape(shape, self._itemsize, order='F')
-                if self.ndim > 1:
-                    self.flags = {'F_CONTIGUOUS': True, 'C_CONTIGUOUS': False}
-                else:
-                   self.flags = {'F_CONTIGUOUS': True, 'C_CONTIGUOUS': True}
-                if self.ndim <= 1:
-                     self.flags = {'F_CONTIGUOUS': True, 'C_CONTIGUOUS': True}
             self._strides = strides
+            self.flags = {
+                'C_CONTIGUOUS': (order == 'C'),
+                'F_CONTIGUOUS': (order == 'F') or (self.ndim <= 1)
+            }
         else:
             # Existing array
             if isinstance(buffer, ndarray) and buffer.base is not None: