feat: implement size_of to make nested field chains work

MrIndeciso · MrIndeciso · commit f8f340823d54 · 2024-10-28T16:43:14.000+01:00
diff --git a/libdestruct/c/__init__.py b/libdestruct/c/__init__.py
@@ -9,5 +9,5 @@
 
 __all__ = ["c_char", "c_uchar", "c_short", "c_ushort", "c_int", "c_uint", "c_long", "c_ulong", "c_str"]
 
-import libdestruct.c.base_type_inflater  # noqa: F401
+import libdestruct.c.base_type_inflater
 import libdestruct.c.ctypes_generic_field  # noqa: F401
diff --git a/libdestruct/common/array/array_impl.py b/libdestruct/common/array/array_impl.py
@@ -8,16 +8,15 @@
 
 from typing import TYPE_CHECKING
 
-from libdestruct.c.ctypes_generic import _ctypes_generic
 from libdestruct.common.array.array import array
-from libdestruct.common.field import Field
-from libdestruct.common.obj import obj
 from libdestruct.common.struct.struct import struct
+from libdestruct.common.utils import size_of
 
 if TYPE_CHECKING: # pragma: no cover
     from collections.abc import Generator
 
     from libdestruct.backing.resolver import Resolver
+    from libdestruct.common.obj import obj
 
 
 class array_impl(array):
@@ -26,6 +25,9 @@ class array_impl(array):
     size: int
     """The size of the array."""
 
+    item_size: int
+    """The size of each item in the array."""
+
     def __init__(
         self: array_impl,
         resolver: Resolver,
@@ -37,15 +39,8 @@ def __init__(
 
         self.backing_type = backing_type
         self._count = count
-
-        if hasattr(backing_type, "size"):
-            self.size = self.backing_type.size * self._count
-            self.item_size = self.backing_type.size
-        elif callable(backing_type):
-            self.size = 8 * self._count
-            self.item_size = 8
-        else:
-            raise NotImplementedError("Unsupported backing type.")
+        self.item_size = size_of(self.backing_type)
+        self.size = self.item_size * self._count
 
     def count(self: array_impl) -> int:
         """Get the size of the array."""
diff --git a/libdestruct/common/array/linear_array_field.py b/libdestruct/common/array/linear_array_field.py
@@ -10,8 +10,9 @@
 
 from libdestruct.common.array.array_field import ArrayField
 from libdestruct.common.array.array_impl import array_impl
+from libdestruct.common.utils import size_of
 
-if TYPE_CHECKING: # pragma: no cover
+if TYPE_CHECKING:  # pragma: no cover
     from libdestruct.backing.resolver import Resolver
     from libdestruct.common.array.array import array
     from libdestruct.common.obj import obj
@@ -32,3 +33,7 @@ def inflate(self: LinearArrayField, resolver: Resolver) -> array:
             resolver: The backing resolver for the object.
         """
         return array_impl(resolver, self.item, self.size)
+
+    def get_size(self: LinearArrayField) -> int:
+        """Returns the size of the object inflated by this field."""
+        return size_of(self.item) * self.size
diff --git a/libdestruct/common/enum/int_enum_field.py b/libdestruct/common/enum/int_enum_field.py
@@ -54,3 +54,7 @@ def inflate(self: IntEnumField, resolver: Resolver) -> int:
             resolver: The backing resolver for the object.
         """
         return enum(resolver, self.enum, self.backing_type, self.lenient)
+
+    def get_size(self: IntEnumField) -> int:
+        """Returns the size of the object inflated by this field."""
+        return self.backing_type.size
diff --git a/libdestruct/common/field.py b/libdestruct/common/field.py
@@ -9,7 +9,7 @@
 from abc import ABC, abstractmethod
 from typing import TYPE_CHECKING
 
-if TYPE_CHECKING: # pragma: no cover
+if TYPE_CHECKING:  # pragma: no cover
     from libdestruct.backing.resolver import Resolver
     from libdestruct.common.obj import obj
 
@@ -26,3 +26,7 @@ def inflate(self: Field, resolver: Resolver) -> obj:
         Args:
             resolver: The backing resolver for the object.
         """
+
+    @abstractmethod
+    def get_size(self: Field) -> int:
+        """Returns the size of the object inflated by this field."""
diff --git a/libdestruct/common/ptr/ptr_factory.py b/libdestruct/common/ptr/ptr_factory.py
@@ -8,18 +8,18 @@
 
 from typing import TYPE_CHECKING
 
-from libdestruct.common.ptr.ptr_struct_field import PtrStructField
+from libdestruct.common.ptr.ptr_field import PtrField
 
 if TYPE_CHECKING: # pragma: no cover
     from libdestruct.common.field import Field
     from libdestruct.common.obj import obj
 
 
-def ptr_to(item: obj | Field) -> PtrStructField:
+def ptr_to(item: obj | Field) -> PtrField:
     """Crafts a struct member which is a pointer to an object."""
-    return PtrStructField(item)
+    return PtrField(item)
 
 
-def ptr_to_self() -> PtrStructField:
+def ptr_to_self() -> PtrField:
     """Crafts a struct member which is a pointer to a struct of the same type."""
-    return PtrStructField(None)
+    return PtrField(None)
diff --git a/libdestruct/common/ptr/ptr_field.py b/libdestruct/common/ptr/ptr_field.py
@@ -10,27 +10,26 @@
 
 from libdestruct.common.field import Field
 from libdestruct.common.ptr.ptr import ptr
-from libdestruct.common.struct.struct_field import StructField
 
 if TYPE_CHECKING: # pragma: no cover
     from libdestruct.backing.resolver import Resolver
     from libdestruct.common.obj import obj
 
 
-class PtrStructField(StructField):
+class PtrField(Field):
     """A generator for a field of a struct."""
 
     base_type: type[obj] = ptr
 
-    def __init__(self: PtrStructField, backing_type: type | Field) -> None:
+    def __init__(self: PtrField, backing_type: type | Field) -> None:
         """Initialize a pointer field.
 
         Args:
             backing_type: The type of the object the pointer points to.
         """
         self.backing_type = backing_type
 
-    def inflate(self: PtrStructField, resolver: Resolver) -> obj:
+    def inflate(self: PtrField, resolver: Resolver) -> obj:
         """Inflate the field.
 
         Args:
@@ -40,3 +39,7 @@ def inflate(self: PtrStructField, resolver: Resolver) -> obj:
             return ptr(resolver, self.backing_type.inflate)
 
         return ptr(resolver, self.backing_type)
+
+    def get_size(self: PtrField) -> int:
+        """Returns the size of the object inflated by this field."""
+        return ptr.size
diff --git a/libdestruct/common/ptr/ptr_field_inflater.py b/libdestruct/common/ptr/ptr_field_inflater.py
@@ -8,7 +8,7 @@
 
 from typing import TYPE_CHECKING
 
-from libdestruct.common.ptr.ptr_struct_field import PtrStructField
+from libdestruct.common.ptr.ptr_field import PtrField
 from libdestruct.common.type_registry import TypeRegistry
 
 if TYPE_CHECKING: # pragma: no cover
@@ -21,7 +21,7 @@
 
 
 def ptr_field_inflater(
-    field: PtrStructField,
+    field: PtrField,
     _: type[obj],
     owner: tuple[obj, type[obj]] | None,
 ) -> Callable[[Resolver], obj]:
@@ -37,4 +37,4 @@ def ptr_field_inflater(
     return field.inflate
 
 
-registry.register_instance_handler(PtrStructField, ptr_field_inflater)
+registry.register_instance_handler(PtrField, ptr_field_inflater)
diff --git a/libdestruct/common/struct/__init__.py b/libdestruct/common/struct/__init__.py
@@ -10,5 +10,5 @@
 
 __all__ = ["struct", "struct_impl", "ptr_to", "ptr_to_self"]
 
-import libdestruct.common.struct.struct_field_inflater
+import libdestruct.common.ptr.ptr_field_inflater
 import libdestruct.common.struct.struct_inflater  # noqa: F401
diff --git a/libdestruct/common/struct/struct_field.py b/libdestruct/common/struct/struct_field.py
diff --git a/libdestruct/common/utils.py b/libdestruct/common/utils.py
@@ -0,0 +1,33 @@
+#
+# This file is part of libdestruct (https://github.com/mrindeciso/libdestruct).
+# Copyright (c) 2024 Roberto Alessandro Bertolini. All rights reserved.
+# Licensed under the MIT license. See LICENSE file in the project root for details.
+#
+
+from __future__ import annotations
+
+from types import MethodType
+from typing import TYPE_CHECKING
+
+from libdestruct.common.field import Field
+
+if TYPE_CHECKING:  # pragma: no cover
+    from libdestruct.backing.resolver import Resolver
+    from libdestruct.common.obj import obj
+
+
+def size_of(item_or_inflater: obj | callable[[Resolver], obj]) -> int:
+    """Return the size of an object, from an obj or it's inflater."""
+    if hasattr(item_or_inflater, "size"):
+        return item_or_inflater.size
+
+    # Check if item is the bound method of a Field
+    if not isinstance(item_or_inflater, MethodType):
+        raise TypeError("Provided inflater is not the bound method of a Field object")
+
+    field_object = item_or_inflater.__self__
+
+    if not isinstance(field_object, Field):
+        raise TypeError("Provided inflater is not the bound method of a Field object")
+
+    return field_object.get_size()
