oh how i love null checks

Author: ZeroIntensity

src/mod.c

@@ -3,8 +3,9 @@
 #include <signal.h>
 #include <setjmp.h>
 #include <stdbool.h>
+#include <stdio.h>
+#include <frameobject.h> // needed to get members of PyFrameObject
 #define GETOBJ PyObject* obj; if (!PyArg_ParseTuple(args, "O", &obj)) return NULL
-#define CALL_ATTR(ob, attr) PyObject_Call(PyObject_GetAttrString(ob, attr), PyTuple_New(0), NULL)
 static jmp_buf buf;
 
 static PyObject* add_ref(PyObject* self, PyObject* args) {
@@ -47,8 +48,7 @@ void handler(int signum) {
 static PyObject* handle(PyObject* self, PyObject* args) {
     PyObject* func;
     PyObject* params = NULL;
-    PyObject* kwargs = NULL;
-    PyObject* faulthandler = PyImport_ImportModule("faulthandler");
+    PyObject* kwargs = NULL;;
 
     if (!PyArg_ParseTuple(
             args,
@@ -64,27 +64,26 @@ static PyObject* handle(PyObject* self, PyObject* args) {
     if (!params) params = PyTuple_New(0);
     if (!kwargs) kwargs = PyDict_New();
 
-    int val = setjmp(buf);
-
-    CALL_ATTR(faulthandler, "disable");
-    // faulthandler needs to be shut off in case of a segfault or its message will still print
-
     if (setjmp(buf)) {
-        CALL_ATTR(faulthandler, "enable");
+        PyFrameObject* frame = PyEval_GetFrame();
+        PyCodeObject* code = frame->f_code;
+        Py_INCREF(code);
+
+        // this is basically a copy of PyFrame_GetCode, which is only available on 3.9+
 
         PyErr_Format(
             PyExc_RuntimeError,
             "segment violation occured during execution of %S",
-            PyObject_GetAttrString(func, "__name__")
+            code->co_name
         );
+
+        Py_DECREF(code);
         return NULL;
     }
 
     PyObject* result = PyObject_Call(func, params, kwargs);
-
     if (!result) return NULL;
 
-    CALL_ATTR(faulthandler, "enable");
     return result;
 }
 

src/pointers/__init__.py

@@ -14,7 +14,7 @@
     TypedCPointer, VoidPointer, array, cast, to_c_ptr, to_struct_ptr, to_voidp
 )
 from .calloc import AllocatedArrayPointer, calloc
-from .constants import NULL, Nullable, raw_type
+from .constants import NULL, Nullable, handle, raw_type
 from .custom_binding import binding, binds
 from .decay import decay, decay_annotated, decay_wrapped
 from .exceptions import (

src/pointers/base_pointers.py

@@ -12,10 +12,10 @@
 
 from typing_extensions import final
 
-from _pointers import add_ref, handle, remove_ref
+from _pointers import add_ref, remove_ref
 
 from ._utils import deref, force_set_attr, move_to_mem
-from .constants import NULL, Nullable
+from .constants import NULL, Nullable, handle
 from .exceptions import DereferenceError, FreedMemoryError, NullPointerError
 
 __all__ = (
@@ -96,22 +96,13 @@ def ensure(self) -> int:
 
 class Movable(ABC, Generic[T, A]):
     @abstractmethod
-    def _move(
-        self,
-        data: Union[A, T],
-        *,
-        unsafe: bool = False,
-    ) -> None:
-        ...
-
-    @final
     def move(
         self,
         data: Union[A, T],
         *,
         unsafe: bool = False,
     ) -> None:
-        handle(self._move, (data,), {"unsafe": unsafe})
+        ...
 
     def __ilshift__(self, data: Union[A, T]):
         self.move(data)
@@ -126,16 +117,12 @@ class Dereferencable(ABC, Generic[T]):
     """Abstract class for an object that may be dereferenced."""
 
     @abstractmethod
-    def _dereference(self) -> T:
-        ...
-
-    @final
     def dereference(self) -> T:
         """Dereference the pointer.
 
         Returns:
             Value at the pointers address."""
-        return handle(self._dereference)
+        ...
 
     @final
     def __invert__(self) -> T:
@@ -204,6 +191,8 @@ def size(self) -> int:
         """Size of the target value."""
         ...
 
+    @handle
+    @final
     def make_ct_pointer(self) -> "ctypes._PointerLike":
         """Convert the address to a ctypes pointer.
 
@@ -262,13 +251,15 @@ def type(self) -> Type[T]:
 
         return self._type
 
+    @handle
     def set_attr(self, key: str, value: Any) -> None:
         v: Any = ~self  # mypy gets angry if this isnt any
         if not isinstance(~self, type):
             v = type(v)
 
         force_set_attr(v, key, value)
 
+    @handle
     def assign(
         self,
         target: Nullable[Union["BaseObjectPointer[T]", T]],
@@ -299,7 +290,7 @@ def assign(
     def address(self) -> Optional[int]:
         return self._address
 
-    def _dereference(self) -> T:
+    def dereference(self) -> T:
         return deref(self.ensure())
 
     def __irshift__(
@@ -366,7 +357,8 @@ def _make_stream_and_ptr(
         bytes_a = (ctypes.c_ubyte * size).from_address(address)
         return self.make_ct_pointer(), bytes(bytes_a)
 
-    def _move(
+    @handle
+    def move(
         self,
         data: Union["BaseCPointer[T]", T],
         *,
@@ -392,6 +384,7 @@ def __ixor__(self, data: Union["BaseCPointer[T]", T]):
         self.move(data, unsafe=True)
         return self
 
+    @handle
     def make_ct_pointer(self):
         return ctypes.cast(
             self.ensure(),
@@ -436,7 +429,8 @@ def assigned(self) -> bool:
     def assigned(self, value: bool) -> None:
         self._assigned = value
 
-    def _move(
+    @handle
+    def move(
         self,
         data: Union[BasePointer[T], T],
         unsafe: bool = False,
@@ -456,7 +450,8 @@ def _move(
         self.assigned = True
         remove_ref(data)
 
-    def _dereference(self) -> T:
+    @handle
+    def dereference(self) -> T:
         if self.freed:
             raise FreedMemoryError(
                 "cannot dereference memory that has been freed",

src/pointers/c_pointer.py

@@ -6,6 +6,7 @@
 
 from ._utils import get_mapped, map_type
 from .base_pointers import BaseCPointer, IterDereferencable, Typed
+from .constants import handle
 
 if TYPE_CHECKING:
     from .structure import Struct, StructPointer
@@ -31,11 +32,13 @@ class VoidPointer(BaseCPointer[Any]):
     def size(self) -> int:
         return self._size
 
-    @property
+    @property  # type: ignore
+    @handle
     def _as_parameter_(self) -> ctypes.c_void_p:
         return ctypes.c_void_p(self.address)
 
-    def _dereference(self) -> Optional[int]:
+    @handle
+    def dereference(self) -> Optional[int]:
         """Dereference the pointer."""
         deref = ctypes.c_void_p.from_address(self.ensure())
         return deref.value
@@ -61,6 +64,7 @@ def decref(self) -> bool:
         """Whether the target objects reference count should be decremented when the pointer is garbage collected."""  # noqa
         ...
 
+    @handle
     def _cleanup(self):
         if self.address:
             if (type(~self) is not str) and (self.decref):
@@ -120,18 +124,20 @@ class TypedCPointer(_TypedPointer[T]):
     def address(self) -> Optional[int]:
         return self._address
 
-    @property
+    @property  # type: ignore
+    @handle
     def _as_parameter_(self):
         ctype = get_mapped(self.type)
         deref = ctype.from_address(self.ensure())
         value = deref.value  # type: ignore
 
         if isinstance(value, (TypedCPointer, VoidPointer)):
-            return ctypes.pointer(value._as_parameter_)
+            return ctypes.pointer(value._as_parameter_)  # type: ignore
 
         return ctypes.pointer(deref)
 
-    def _dereference(self) -> T:
+    @handle
+    def dereference(self) -> T:
         """Dereference the pointer."""
         ctype = get_mapped(self.type)
 
@@ -177,14 +183,16 @@ def decref(self) -> bool:
     def type(self) -> Type[T]:  # type: ignore
         return self._type
 
-    @property
+    @property  # type: ignore
+    @handle
     def _as_parameter_(self) -> "ctypes.Array[ctypes._CData]":
         ctype = get_mapped(self.type)
 
         deref = (ctype * self._length).from_address(self.ensure())
         return deref
 
-    def _dereference(self) -> List[T]:
+    @handle
+    def dereference(self) -> List[T]:
         """Dereference the pointer."""
         array = self._as_parameter_
         return [array[i] for i in range(self._length)]  # type: ignore
@@ -200,6 +208,7 @@ def __getitem__(self, index: int) -> T:
         return array[index]
 
 
+@handle
 def cast(ptr: VoidPointer, data_type: Type[T]) -> TypedCPointer[T]:
     """Cast a void pointer to a typed pointer."""
 
@@ -237,6 +246,7 @@ def to_struct_ptr(struct: A) -> "StructPointer[A]":
     return StructPointer(id(struct), type(struct))
 
 
+@handle
 def array(*seq: T) -> CArrayPointer[T]:
     f_type = type(seq[0])
 

src/pointers/calloc.py

@@ -2,6 +2,7 @@
 
 from ._cstd import c_calloc, c_free
 from .base_pointers import BaseAllocatedPointer
+from .constants import handle
 from .exceptions import AllocationError
 
 __all__ = ("AllocatedArrayPointer", "calloc")
@@ -108,6 +109,7 @@ def __setitem__(self, index: int, value: T) -> None:
         chunk = self._get_chunk_at(index)
         chunk <<= value
 
+    @handle
     def free(self) -> None:
         first = self[0]
         first.ensure_valid()

src/pointers/constants.py

@@ -1,13 +1,30 @@
 import ctypes
-from typing import Any, NamedTuple, Type, TypeVar, Union
+import faulthandler
+from contextlib import suppress
+from functools import wraps
+from io import UnsupportedOperation
+from typing import Any, Callable, NamedTuple, Type, TypeVar, Union
+
+from typing_extensions import ParamSpec
+
+from _pointers import handle as _handle
+
+from .exceptions import SegmentViolation
+
+with suppress(
+    UnsupportedOperation
+):  # in case its running in idle or something like that
+    faulthandler.enable()
 
 __all__ = (
     "NULL",
     "Nullable",
     "raw_type",
+    "handle",
 )
 
 T = TypeVar("T")
+P = ParamSpec("P")
 
 
 class NULL:
@@ -24,3 +41,28 @@ class RawType(NamedTuple):
 def raw_type(ct: Type["ctypes._CData"]) -> Any:
     """Set a raw ctypes type for a struct."""
     return RawType(ct)
+
+
+def handle(func: Callable[P, T]) -> Callable[P, T]:
+    @wraps(func)
+    def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
+        try:
+            faulthandler.disable()
+            call = _handle(func, args, kwargs)
+
+            with suppress(UnsupportedOperation):
+                faulthandler.enable()
+
+            return call
+        except RuntimeError as e:
+            msg = str(e)
+
+            if not msg.startswith("segment"):
+                raise
+
+            with suppress(UnsupportedOperation):
+                faulthandler.enable()
+
+            raise SegmentViolation(str(e)) from None
+
+    return wrapper

src/pointers/exceptions.py

@@ -6,6 +6,7 @@
     "InvalidBindingParameter",
     "NullPointerError",
     "InvalidVersionError",
+    "SegmentViolation",
 )
 
 
@@ -49,3 +50,9 @@ class InvalidVersionError(Exception):
     """Python version is not high enough."""
 
     pass
+
+
+class SegmentViolation(Exception):
+    """SIGSEGV was sent to Python."""
+
+    pass