probe: improved *probe class, functions & test

src/liburing/probe.pxd

@@ -1,11 +1,13 @@
-from cpython.mem cimport PyMem_RawCalloc
+from cpython.mem cimport PyMem_RawCalloc, PyMem_RawFree
 from .lib.uring cimport *
 from .error cimport trap_error, memory_error
 from .queue cimport io_uring
 
 
 cdef class io_uring_probe:
-    cdef __io_uring_probe *ptr
+    cdef:
+        __io_uring_probe*   ptr
+        unsigned int        len
 
 cpdef io_uring_probe io_uring_get_probe_ring(io_uring ring)
 cpdef io_uring_probe io_uring_get_probe()

src/liburing/probe.py

@@ -17,8 +17,10 @@ def probe() -> dict:
     '''
     r = {}
     p = io_uring_get_probe()
-    for i in io_uring_op:
-        if i.name != 'IORING_OP_LAST':
-            r[i.name] = io_uring_opcode_supported(p, i)
-    io_uring_free_probe(p)
+    try:
+        for i in io_uring_op:
+            if i.name != 'IORING_OP_LAST':
+                r[i.name] = io_uring_opcode_supported(p, i)
+    finally:
+        io_uring_free_probe(p)
     return r

src/liburing/probe.pyx

@@ -2,25 +2,29 @@ cdef class io_uring_probe:
 
     def __cinit__(self, unsigned int num=0):
         if num:
-            self.ptr = <__io_uring_probe*>PyMem_RawCalloc(num, sizeof(__io_uring_probe))
+            self.ptr = <__io_uring_probe*>PyMem_RawCalloc(
+                num, sizeof(__io_uring_probe) + num * sizeof(__io_uring_probe_op)
+            )
             if self.ptr is NULL:
                 memory_error(self)
+            self.len = num
 
     def __dealloc__(self):
-        if self.ptr is not NULL:
-            # just in case user forgets to call `io_uring_free_probe` or error happened
-            __io_uring_free_probe(self.ptr)
+        if self.len and self.ptr is not NULL:
+            PyMem_RawFree(self.ptr)
             self.ptr = NULL
 
     @property
     def last_op(self):
         if self.ptr is not NULL:
             return self.ptr.last_op
+        memory_error(self)
 
     @property
     def ops_len(self):
         if self.ptr is not NULL:
             return self.ptr.ops_len
+        memory_error(self)
 
 cpdef io_uring_probe io_uring_get_probe_ring(io_uring ring):
     cdef io_uring_probe probe = io_uring_probe()

test/probe_test.py

@@ -1,3 +1,4 @@
+from re import escape
 from errno import EINVAL
 from pytest import raises
 from liburing import probe, io_uring_get_probe, io_uring_get_probe_ring, io_uring_free_probe, \
@@ -9,17 +10,17 @@ def test_probe():
     assert op['IORING_OP_NOP'] is True
     assert op.get('IORING_OP_LAST', None) is None
 
-    # triggers `__dealloc__`
-    io_uring_get_probe()
+    p = io_uring_get_probe()
+    io_uring_free_probe(p)
+    with raises(MemoryError, match=escape('`io_uring_probe()` is out of memory!')):
+        p.ops_len
 
 
 def test_probe_ring(ring):
     p = io_uring_get_probe_ring(ring)
-
     assert p.ops_len > 31
     assert p.last_op > 30
     assert p.last_op == p.ops_len - 1
-
     io_uring_free_probe(p)
 
 
@@ -28,8 +29,8 @@ def test_probe_register(ring):
     with raises(OSError) as e:
         io_uring_register_probe(ring, p, 256)
     assert e.value.errno == EINVAL
-    io_uring_free_probe(p)
+    # free is doen by the `__dealloc__` since `num` is set
 
     p = io_uring_probe(4)
     assert io_uring_register_probe(ring, p, 4) == 0
-    io_uring_free_probe(p)
+    # free is doen by the `__dealloc__` since `num` is set