Add suggestions

dissect/cstruct/types/base.py

@@ -287,7 +287,7 @@ def _write(cls, stream: BinaryIO, data: list[Any]) -> int:
         return cls.type._write_array(stream, data)
 
     @classmethod
-    def _type_stub(cls, name: str = ""):
+    def _type_stub(cls, name: str = "") -> str:
         return f"{name}: {cls.__base__.__name__}[{cls.type.__name__}]"
 
 

dissect/cstruct/types/packed.py

@@ -65,6 +65,6 @@ def _write_array(cls, stream: BinaryIO, data: list[Packed]) -> int:
         return stream.write(_struct(cls.cs.endian, f"{len(data)}{cls.packchar}").pack(*data))
 
     @classmethod
-    def to_type_stub(cls, name: str):
+    def to_type_stub(cls, name: str) -> str:
         types = ", ".join([x.__name__ for x in cls.__bases__])
-        return f"{name}= type[{types}]"
+        return f"{name} = type[{types}]"

dissect/cstruct/types/structure.py

@@ -366,23 +366,26 @@ def commit(cls) -> None:
         for key, value in classdict.items():
             setattr(cls, key, value)
 
-    def to_type_stub(cls, name: str = ""):
-        with io.StringIO() as data:
-            data.write(f"class {cls.__name__}({cls.__base__.__name__}):\n")
-            call_args = ["self"]
-            for key, field in cls.lookup.items():
-                if isinstance(field.type, StructureMetaType):
-                    class_info = field.type.to_type_stub()
-                    data.write(indent(class_info, prefix=" " * 4))
-                call_args.append(f"{field.type_stub()}=...")
-
-            for field in cls.fields.values():
-                type_info = field.type_stub()
-                data.write(indent(f"{type_info}\n", prefix=" " * 4))
-
-            call = ", ".join(call_args)
-            data.write(indent(f"def __init__({call}): ...\n", prefix=" " * 4))
-            return data.getvalue()
+    def to_type_stub(cls, name: str = "") -> str:
+        buffer = io.StringIO()
+        buffer.write(f"class {cls.__name__}({cls.__base__.__name__}):\n")
+        call_args = ["self"]
+        for key, field in cls.lookup.items():
+            if isinstance(field.type, StructureMetaType):
+                class_info = field.type.to_type_stub()
+                buffer.write(indent(class_info, prefix=" " * 4))
+            call_args.append(f"{field.type_stub()}=...")
+
+        for field in cls.fields.values():
+            type_info = field.type_stub()
+            buffer.write(indent(f"{type_info}\n", prefix=" " * 4))
+
+        call = ", ".join(call_args)
+        buffer.write(indent(f"def __init__({call}): ...\n", prefix=" " * 4))
+
+        output = buffer.getvalue()
+        buffer.close()
+        return output
 
 
 class Structure(BaseType, metaclass=StructureMetaType):