Improve partial structs containing self-references

Author: monkeyman192

docs/docs/gui/hex_view.rst

@@ -10,7 +10,7 @@ Note that the view can only access memory that the process has access to, so ent
 Methods
 -------
 
-The ``gui`` property of the mod is a :py:class:`~pymhf.gui.gui.GUI` instance which has the ``hex_view`` property.
+The ``pymhf_gui`` property of the mod is a :py:class:`~pymhf.gui.gui.GUI` instance which has the ``hex_view`` property.
 
 This :py:class:`~pymhf.gui.hexview.HexView` instance has two methods which are useful:
 

pymhf/core/hooking.py

@@ -695,7 +695,15 @@ def call_custom_callbacks(
             if kwargs is None:
                 kwargs = {}
             for cb in callbacks.get(detour_time, set()):
-                cb(*args, **kwargs)
+                try:
+                    cb(*args, **kwargs)
+                except Exception:
+                    logger.exception(f"There was an issue calling custom callback {cb}. It has been removed.")
+                    self._remove_custom_callbacks(
+                        {
+                            cb,
+                        }
+                    )
         elif alert_nonexist:
             raise ValueError(f"Custom callback {callback_key} cannot be found.")
 

pymhf/core/mod_loader.py

@@ -189,6 +189,7 @@ class Mod(ABC):
     _disabled: bool = False
 
     def __init__(self):
+        self._abc_initialised = True
         # Find all the hooks defined for the mod.
         self.hooks: set[HookProtocol] = self.get_members(_funchook_predicate)
         self._custom_callbacks = self.get_members(_callback_predicate)
@@ -199,6 +200,7 @@ def __init__(self):
         self._gui_comboboxes: dict[str, ComboBoxProtocol] = {
             x[1].__qualname__: x[1] for x in inspect.getmembers(self, _gui_combobox_predicate)
         }
+        # TODO: If this isn't initialised and a call is made to it before it is we have an issue...
         self.pymhf_gui = None
         # For variables, unless there is a better way, store just the name so we
         # can our own special binding of the name to the GUI.
@@ -399,9 +401,16 @@ def load_mod_folder(self, folder: str, bind: bool = True, deep_search: bool = Fa
 
         return loaded_mods, bound_hooks
 
-    def instantiate_mod(self, mod: type[Mod], quiet: bool = False) -> Mod:
+    def instantiate_mod(self, mod: type[Mod], quiet: bool = False) -> Optional[Mod]:
         """Register all the functions within the mod as hooks."""
         _mod = mod()
+        # Detect whether or not the mod has called __init__ on the parent class.
+        if not getattr(_mod, "_abc_initialised", False):
+            logger.error(
+                f"The mod {mod} has an __init__ statement which doesn't call super().__init__\n"
+                "This mod will not be loaded until this is fixed."
+            )
+            return None
         # First register each of the methods which are detours.
         for hook in _mod.hooks:
             self.hook_manager.register_hook(hook)
@@ -482,7 +491,9 @@ def reload(self, name: str, gui: "GUI"):
                 new_module = self.load_mod(module.__file__)
                 for _mod in self._preloaded_mods.values():
                     mod = self.instantiate_mod(_mod)
-
+                    if mod is None:
+                        # If the mod isn't instantiated for any reason, skip it.
+                        continue
                     # Get the mod states for the mod if there are any and reapply them to the new mod
                     # instance.
                     if mod_state := self.mod_states.get(name):

pymhf/gui/gui.py

@@ -476,11 +476,15 @@ def run(self):
             while dpg.is_dearpygui_running():
                 # For each tracking variable, update the value.
                 for vars in self.tracking_variables.get(self._current_tab, []):
-                    tag, cls, var, is_str = vars
-                    if is_str:
-                        dpg.set_value(tag, str(getattr(cls, var)))
-                    else:
-                        dpg.set_value(tag, getattr(cls, var))
+                    try:
+                        tag, cls, var, is_str = vars
+                        if is_str:
+                            dpg.set_value(tag, str(getattr(cls, var)))
+                        else:
+                            dpg.set_value(tag, getattr(cls, var))
+                    except Exception:
+                        # If we can't set the value, don't crash the whole program.
+                        pass
                 dpg.render_dearpygui_frame()
             dpg.destroy_context()
         except Exception:

pymhf/utils/partial_struct.py

@@ -40,6 +40,10 @@ def partial_struct(cls: _T) -> _T:
         calling_frame = cframe.f_back
     if calling_frame is not None:
         _locals.update(calling_frame.f_locals)
+
+    # Also add the class to the locals so that it can find itself in the case it references itself.
+    _locals.update({cls.__name__: cls})
+
     _fields_ = []
     curr_position = 0
     total_size = getattr(cls, "_total_size_", 0)
@@ -78,6 +82,17 @@ def partial_struct(cls: _T) -> _T:
             else:
                 field_type = annotation.__origin__
                 field_offset = metadata
+        # Check to make sure that the `field_type` is not a string(as can happen when we have a
+        # self-reference).
+        if isinstance(field_type, str):
+            field_type = getattr(annotation, "__origin__", None)
+            if field_type is None:
+                raise ValueError(
+                    f"The provided metadata {metadata} for field {field_name!r} is invalid. "
+                    "If the type in the `Field` component of the annotation is a string, please ensure the "
+                    "first argument of the Annotation is also the same string so that the type can be "
+                    "resolved."
+                )
         if not issubclass(field_type, get_args(CTYPES)):
             raise ValueError(f"The field {field_name!r} has an invalid type: {field_type}")
         if field_offset is not None and not isinstance(field_offset, int):

tests/test_partial_structs.py

@@ -1,7 +1,8 @@
 import ctypes
 import re
+import types
 from enum import IntEnum
-from typing import Annotated
+from typing import Annotated, Generic, Type, TypeVar, Union
 
 import pytest
 
@@ -405,3 +406,61 @@ class Test4(ctypes.Structure):
 def test_invalid_c_enum32_cases():
     with pytest.raises(TypeError):
         c_enum32[22]
+
+
+T = TypeVar("T", bound=Union[ctypes._SimpleCData, ctypes.Structure])
+
+
+class cTkDynamicArray(ctypes.Structure, Generic[T]):
+    _template_type: Type[T]
+    _fields_ = [
+        ("offset", ctypes.c_uint32),
+        ("count", ctypes.c_uint32),
+    ]
+
+    offset: int
+    count: int
+
+    def value(self, source: bytearray) -> ctypes.Array[T]:
+        # This is pretty hacky, but it does the job.
+        # A more realistic implementation would be reading memory directly so would be implemented as a
+        # property.
+        if self.offset == 0 or self.count == 0:
+            # Empty lists are stored as empty header bytes.
+            return (self._template_type * 0)()
+        type_ = self._template_type * self.count
+        return type_.from_buffer(source, self.offset)
+
+    def __class_getitem__(cls: type["cTkDynamicArray"], key: T):
+        _cls: type["cTkDynamicArray"] = types.new_class(f"cTkDynamicArray<{key}>", (cls,))
+        _cls._template_type = key
+        return _cls
+
+
+def test_self_referential_struct():
+    # Test the case of the struct having a data type which is itself.
+    # To do this we'll need to introduce a serializable list.
+    @partial_struct
+    class SelfRef(ctypes.Structure):
+        a: Annotated[ctypes.c_uint32, 0x0]
+        children: Annotated["cTkDynamicArray[SelfRef]", 0x4]
+
+    data = bytearray(
+        b"\x01\x00\x00\x00"  # 'a' for the parent.
+        b"\x0c\x00\x00\x00\x02\x00\x00\x00"  # Child data "header"
+        b"\x02\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00"  # Child 1. a = 2
+        b"\x03\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00"  # Child 2. a = 3
+    )
+
+    obj = SelfRef.from_buffer(data)
+    assert obj.a == 1
+    assert obj.children.count == 2
+    # Get the children
+    children = obj.children.value(data)
+    assert children[0].a == 2
+    assert children[0].children.count == 0
+    sub_child = children[0].children.value(data)
+    assert sub_child._length_ == 0
+    assert sub_child._type_ == SelfRef
+    assert children[1].a == 3
+    assert children[1].children.count == 0