Add context manager

Author: Schamper

dissect/cstruct/cstruct.py

@@ -7,6 +7,7 @@
 import sys
 import types
 import warnings
+from contextlib import contextmanager
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, BinaryIO, Literal, TypeVar, cast
 
@@ -31,7 +32,7 @@
 from dissect.cstruct.types.base import MetaType, normalize_endianness
 
 if TYPE_CHECKING:
-    from collections.abc import Iterable
+    from collections.abc import Iterable, Iterator
     from typing import TypeAlias
 
     from dissect.cstruct.types import (
@@ -222,9 +223,11 @@ def __copy__(self) -> cstruct:
 
         # Update typedefs to point to the new cstruct instance
         for name, type in self.typedefs.items():
+            # Skip all the default types
             if name in cs.typedefs:
                 continue
 
+            # Update the cstruct instance for all other types
             if isinstance(type, MetaType):
                 new_type = copy.copy(type)
                 new_type.cs = cs
@@ -389,6 +392,25 @@ def copy(self) -> cstruct:
         """
         return copy.copy(self)
 
+    @contextmanager
+    def endianness(self, endian: AllowedEndianness) -> Iterator[cstruct]:
+        """Context manager for temporarily changing the endianness of this cstruct instance."""
+        new = self.copy()
+        new.endian = normalize_endianness(endian)
+        yield new
+
+    @contextmanager
+    def big_endian(self) -> Iterator[cstruct]:
+        """Context manager for temporarily changing the endianness to big endian."""
+        with self.endianness(">") as cs:
+            yield cs
+
+    @contextmanager
+    def little_endian(self) -> Iterator[cstruct]:
+        """Context manager for temporarily changing the endianness to little endian."""
+        with self.endianness("<") as cs:
+            yield cs
+
     def _make_type(
         self,
         name: str,

tests/test_basic.py

@@ -563,3 +563,26 @@ def test_copy(cs: cstruct) -> None:
     # 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
+
+
+def test_endianness(cs: cstruct) -> None:
+    """Test that the endianness context manager works correctly."""
+    with cs.endianness(">") as big_endian_cs:
+        assert big_endian_cs.endian == ">"
+        assert big_endian_cs.uint32(b"\x00\x00\x00\x01") == 1
+
+        # Verify original cstruct is unchanged
+        assert cs.endian == "<"
+        assert cs.uint32(b"\x01\x00\x00\x00") == 1
+
+    with cs.endianness("<") as little_endian_cs:
+        assert little_endian_cs.endian == "<"
+        assert little_endian_cs.uint32(b"\x01\x00\x00\x00") == 1
+
+    with cs.big_endian() as big_endian_cs:
+        assert big_endian_cs.endian == ">"
+        assert big_endian_cs.uint32(b"\x00\x00\x00\x01") == 1
+
+    with cs.little_endian() as little_endian_cs:
+        assert little_endian_cs.endian == "<"
+        assert little_endian_cs.uint32(b"\x01\x00\x00\x00") == 1