Move all internal attributes to dunder attributes

Author: Schamper

dissect/cstruct/bitbuffer.py

@@ -19,11 +19,11 @@ def __init__(self, stream: BinaryIO, *, endian: str):
 
     def read(self, field_type: type[BaseType], bits: int) -> int:
         if self._remaining == 0 or self._type != field_type:
-            if field_type.size is None:
+            if field_type.__size__ is None:
                 raise ValueError("Reading variable-length fields is unsupported")
 
             self._type = field_type
-            self._remaining = field_type.size * 8
+            self._remaining = field_type.__size__ * 8
             self._buffer = field_type._read(self.stream, endian=self.endian)
 
         if isinstance(self._buffer, bytes):
@@ -51,17 +51,17 @@ def write(self, field_type: type[BaseType], data: int, bits: int) -> None:
             if self._type:
                 self.flush()
 
-            if field_type.size is None:
+            if field_type.__size__ is None:
                 raise ValueError("Writing variable-length fields is unsupported")
 
-            self._remaining = field_type.size * 8
+            self._remaining = field_type.__size__ * 8
             self._type = field_type
 
-        if self._type is None or self._type.size is None:
+        if self._type is None or self._type.__size__ is None:
             raise ValueError("Invalid state")
 
         if self.endian == "<":
-            self._buffer |= data << (self._type.size * 8 - self._remaining)
+            self._buffer |= data << (self._type.__size__ * 8 - self._remaining)
         else:
             self._buffer |= data << (self._remaining - bits)
 

dissect/cstruct/compiler.py

@@ -60,7 +60,7 @@
 
 
 def compile(structure: type[Structure]) -> type[Structure]:
-    return Compiler(structure.cs).compile(structure)
+    return Compiler(structure.__cs__).compile(structure)
 
 
 class Compiler:
@@ -119,7 +119,7 @@ def generate_source(self) -> str:
         """
 
         if any(field.bits for field in self.fields):
-            preamble += "bit_reader = BitBuffer(stream, endian=endian)\n"
+            preamble += "bit_reader = BitBuffer(stream, endian=cls.__cs__.endian)\n"
 
         read_code = "\n".join(self._generate_fields())
 
@@ -224,22 +224,22 @@ def align_to_field(field: Field) -> Iterator[str]:
         yield from flush()
 
         if self.align:
-            yield f"stream.seek(-stream.tell() & (cls.alignment - 1), {io.SEEK_CUR})"
+            yield f"stream.seek(-stream.tell() & (cls.__alignment__ - 1), {io.SEEK_CUR})"
 
     def _generate_structure(self, field: Field) -> Iterator[str]:
         template = f"""
-        {"_s = stream.tell()" if field.type.dynamic else ""}
+        {"_s = stream.tell()" if field.type.__dynamic__ else ""}
         r["{field._name}"] = {self._map_field(field)}._read(stream, context=r, endian=endian)
-        {f's["{field._name}"] = stream.tell() - _s' if field.type.dynamic else ""}
+        {f's["{field._name}"] = stream.tell() - _s' if field.type.__dynamic__ else ""}
         """
 
         yield dedent(template)
 
     def _generate_array(self, field: Field) -> Iterator[str]:
         template = f"""
-        {"_s = stream.tell()" if field.type.dynamic else ""}
+        {"_s = stream.tell()" if field.type.__dynamic__ else ""}
         r["{field._name}"] = {self._map_field(field)}._read(stream, context=r, endian=endian)
-        {f's["{field._name}"] = stream.tell() - _s' if field.type.dynamic else ""}
+        {f's["{field._name}"] = stream.tell() - _s' if field.type.__dynamic__ else ""}
         """
 
         yield dedent(template)
@@ -253,8 +253,8 @@ def _generate_bits(self, field: Field) -> Iterator[str]:
             field_type = field_type.type
 
         if issubclass(field_type, Char):
-            field_type = field_type.cs.uint8
-            lookup = "cls.cs.uint8"
+            field_type = field_type.__cs__.uint8
+            lookup = "cls.__cs__.uint8"
 
         template = f"""
         _t = {lookup}
@@ -283,13 +283,13 @@ def _generate_packed(self, fields: list[Field]) -> Iterator[str]:
                 read_type = _get_read_type(self.cs, field_type.type)
 
                 if issubclass(read_type, (Char, Wchar, Int)):
-                    count *= read_type.size
+                    count *= read_type.__size__
                     getter = f"buf[{size}:{size + count}]"
                 else:
                     getter = f"data[{slice_index}:{slice_index + count}]"
                     slice_index += count
             elif issubclass(read_type, (Char, Wchar, Int)):
-                getter = f"buf[{size}:{size + read_type.size}]"
+                getter = f"buf[{size}:{size + read_type.__size__}]"
             else:
                 getter = f"data[{slice_index}]"
                 slice_index += 1
@@ -308,8 +308,8 @@ def _generate_packed(self, fields: list[Field]) -> Iterator[str]:
 
                 if issubclass(field_type.type, Int):
                     reads.append(f"_b = {getter}")
-                    item_parser = parser_template.format(type="_et", getter=f"_b[i:i + {field_type.type.size}]")
-                    list_comp = f"[{item_parser} for i in range(0, {count}, {field_type.type.size})]"
+                    item_parser = parser_template.format(type="_et", getter=f"_b[i:i + {field_type.type.__size__}]")
+                    list_comp = f"[{item_parser} for i in range(0, {count}, {field_type.type.__size__})]"
                 elif issubclass(field_type.type, Pointer):
                     item_parser = "_et.__new__(_et, e, stream, context=r, endian=endian)"
                     list_comp = f"[{item_parser} for e in {getter}]"
@@ -329,7 +329,7 @@ def _generate_packed(self, fields: list[Field]) -> Iterator[str]:
             reads.append(f'r["{field._name}"] = {parser}')
             reads.append("")  # Generates a newline in the resulting code
 
-            size += field_type.size
+            size += field_type.__size__
 
         fmt = _optimize_struct_fmt(info)
         if fmt == "x" or (len(fmt) == 2 and fmt[1] == "x"):
@@ -382,9 +382,9 @@ def _generate_struct_info(cs: cstruct, fields: list[Field], align: bool = False)
         # Other types are byte based
         # We don't actually unpack anything here but slice directly out of the buffer
         elif issubclass(read_type, (Char, Wchar, Int)):
-            yield field, count * read_type.size, "x"
+            yield field, count * read_type.__size__, "x"
 
-        size = count * read_type.size
+        size = count * read_type.__size__
         imaginary_offset += size
         if current_offset is not None:
             current_offset += size

dissect/cstruct/cstruct.py

@@ -206,7 +206,28 @@ def __init__(self, load: str = "", *, endian: AllowedEndianness = "<", pointer:
         if load:
             self.load(load)
 
-    def __getattr__(self, attr: str) -> Any:
+    def __getattr__(
+        self, attr: str
+    ) -> (
+        type[
+            LEB128
+            | BaseType
+            | Char
+            | Enum
+            | Flag
+            | Int
+            | Packed[int]
+            | Packed[float]
+            | Pointer
+            | Structure
+            | Union
+            | Void
+            | Wchar
+        ]
+        | int
+        | str
+        | bytes
+    ):
         try:
             return self.consts[attr]
         except KeyError:
@@ -227,7 +248,7 @@ def __copy__(self) -> cstruct:
         # Update types to point to the new cstruct instance
         for name, type_ in self.types.items():
             new_type = copy.copy(type_)
-            new_type.cs = cs
+            new_type.__cs__ = cs
             cs.add_type(name, new_type, replace=True)
 
         for name, value in self.consts.items():
@@ -481,10 +502,10 @@ def _make_type(
         attrs = attrs or {}
         attrs.update(
             {
-                "cs": self,
-                "size": size,
-                "dynamic": size is None,
-                "alignment": alignment or size,
+                "__cs__": self,
+                "__size__": size,
+                "__dynamic__": size is None,
+                "__alignment__": alignment or size,
             }
         )
         return types.new_class(name, bases, {}, lambda ns: ns.update(attrs))
@@ -494,12 +515,12 @@ def _make_array(self, type_: T, num_entries: int | Expression | None) -> type[Ar
         if num_entries is None:
             null_terminated = True
             size = None
-        elif isinstance(num_entries, Expression) or type_.dynamic:
+        elif isinstance(num_entries, Expression) or type_.__dynamic__:
             size = None
         else:
-            if type_.size is None:
+            if type_.__size__ is None:
                 raise ValueError(f"Cannot create array of dynamic type: {type_.__name__}")
-            size = num_entries * type_.size
+            size = num_entries * type_.__size__
 
         name = f"{type_.__name__}[]" if null_terminated else f"{type_.__name__}[{num_entries}]"
 
@@ -511,7 +532,7 @@ def _make_array(self, type_: T, num_entries: int | Expression | None) -> type[Ar
             "null_terminated": null_terminated,
         }
 
-        return cast("type[Array]", self._make_type(name, bases, size, alignment=type_.alignment, attrs=attrs))
+        return cast("type[Array]", self._make_type(name, bases, size, alignment=type_.__alignment__, attrs=attrs))
 
     def _make_int_type(self, name: str, size: int, signed: bool, *, alignment: int | None = None) -> type[Int]:
         return cast("type[Int]", self._make_type(name, (Int,), size, alignment=alignment, attrs={"signed": signed}))
@@ -538,8 +559,8 @@ def _make_pointer(self, target: type[BaseType]) -> type[Pointer]:
         return self._make_type(
             f"{target.__name__}*",
             (Pointer,),
-            self.pointer.size,
-            alignment=self.pointer.alignment,
+            self.pointer.__size__,
+            alignment=self.pointer.__alignment__,
             attrs={"type": target},
         )
 

dissect/cstruct/tools/stubgen.py

@@ -159,15 +159,15 @@ def generate_structure_stub(
     indent = " " * 4
 
     args = ["self"]
-    for field_name, field in structure.fields.items():
+    for field_name, field in structure.__members__.items():
         inlined = False
 
         # If it's a structure and not globally defined, add an inline stub for it
         nested_type = field.type
         while issubclass(nested_type, types.BaseArray):
             nested_type = nested_type.type
 
-        if issubclass(nested_type, types.Structure) and nested_type.__name__ not in structure.cs.types:
+        if issubclass(nested_type, types.Structure) and nested_type.__name__ not in structure.__cs__.types:
             inlined = True
             inline_stub = generate_structure_stub(nested_type, cs_prefix=cs_prefix, module_prefix=module_prefix)