Add copy method to cstruct instance (#140)

Author: Schamper

dissect/cstruct/cstruct.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import copy
 import ctypes as _ctypes
 import struct
 import sys
@@ -25,6 +26,7 @@
     Void,
     Wchar,
 )
+from dissect.cstruct.types.base import MetaType
 
 if TYPE_CHECKING:
     from collections.abc import Iterable
@@ -206,6 +208,25 @@ def __getattr__(self, attr: str) -> Any:
 
         raise AttributeError(f"Invalid attribute: {attr}")
 
+    def __copy__(self) -> cstruct:
+        cs = cstruct(endian=self.endian, pointer=self.pointer.__name__)
+        cs._anonymous_count = self._anonymous_count
+        cs.consts = self.consts.copy()
+        cs.lookups = self.lookups.copy()
+        cs.includes = self.includes.copy()
+
+        # Update typedefs to point to the new cstruct instance
+        for name, type in self.typedefs.items():
+            if name in cs.typedefs:
+                continue
+
+            if isinstance(type, MetaType):
+                new_type = copy.copy(type)
+                new_type.cs = cs
+                cs.typedefs[name] = new_type
+
+        return cs
+
     def _next_anonymous(self) -> str:
         name = f"__anonymous_{self._anonymous_count}__"
         self._anonymous_count += 1
@@ -328,6 +349,14 @@ def resolve(self, name: type[BaseType] | str) -> type[BaseType]:
 
         raise ResolveError(f"Recursion limit exceeded while resolving type {name}")
 
+    def copy(self) -> cstruct:
+        """Create a copy of this cstruct instance.
+
+        Returns:
+            A new cstruct instance with the same types and settings as this one.
+        """
+        return copy.copy(self)
+
     def _make_type(
         self,
         name: str,

tests/test_basic.py

@@ -532,3 +532,32 @@ def test_linked_list(cs: cstruct) -> None:
 
     assert obj.data == 1
     assert obj.next.data == 2
+
+
+def test_copy(cs: cstruct) -> None:
+    """Test that copying a cstruct instance works as expected."""
+    cdef = """
+    struct test {
+        uint32  a;
+    };
+    """
+    cs.load(cdef)
+
+    cs_copy = cs.copy()
+
+    # Modify the copy and verify the original is unchanged
+    cs_copy.load("""
+    struct test2 {
+        uint16  b;
+    };
+    """)
+
+    assert "test" in cs.typedefs
+    assert "test2" not in cs.typedefs
+
+    assert "test" in cs_copy.typedefs
+    assert "test2" in cs_copy.typedefs
+
+    # Verify that types in the copied cstruct reference the copied cstruct
+    assert cs_copy.resolve("test").cs is cs_copy
+    assert cs_copy.resolve("test2").cs is cs_copy