Merge pull request #90 from ArcanaFramework/touch-ups

Author: tclose

extras/fileformats/extras/application/medical.py

@@ -1,12 +1,17 @@
+import sys
 import typing as ty
 from pathlib import Path
-from typing_extensions import TypeAlias
 import pydicom.tag
 from fileformats.core import FileSet, extra_implementation
 from fileformats.application import Dicom
 import medimages4tests.dummy.dicom.mri.t1w.siemens.skyra.syngo_d13c
 from fileformats.core import SampleFileGenerator
 
+if sys.version_info <= (3, 11):
+    from typing_extensions import TypeAlias
+else:
+    from typing import TypeAlias
+
 TagListType: TypeAlias = ty.Union[
     ty.List[int],
     ty.List[str],

fileformats/core/sampling.py

@@ -6,6 +6,9 @@
 import fileformats.core
 
 
+FILE_FILL_LENGTH_DEFAULT = 256
+
+
 class SampleFileGenerator:
     """Generates sample files. Designed to be used within generate_sample_data overrides
 
@@ -48,6 +51,14 @@ def _generate_fname_stem(self) -> str:
     def rng(self) -> random.Random:
         return random.Random(self.seed)
 
+    def generate_contents(
+        self, binary: bool, fill: int = FILE_FILL_LENGTH_DEFAULT
+    ) -> ty.Union[str, bytes]:
+        if binary:
+            return bytes(self.rng.choices(list(range(256)), k=fill))
+        else:
+            return "".join(self.rng.choices(string.printable, k=fill))
+
     def generate(
         self,
         file_type: ty.Union[
@@ -88,14 +99,12 @@ def generate(
         fspath.parent.mkdir(parents=True, exist_ok=True)
         is_binary: bool = getattr(file_type, "binary", False)
         if not contents:
-            contents = (
-                bytes(random.choices(list(range(256)), k=fill))
-                if is_binary
-                else "".join(random.choices(string.printable, k=fill))
+            contents = self.generate_contents(
+                is_binary, fill if fill else FILE_FILL_LENGTH_DEFAULT
             )
         else:
             contents_type = bytes if is_binary else str
-            if not isinstance(contents, bytes):
+            if not isinstance(contents, contents_type):
                 raise TypeError(
                     f"contents must be {contents_type} for {file_type} files, "
                     f"not {type(contents)}"

fileformats/generic/generate_sample_data.py

@@ -30,12 +30,19 @@ def file_generate_sample_data(
     contents = None
     if getattr(file, "binary", False):
         if hasattr(file, "magic_number"):
+            contents = generator.generate_contents(binary=True)
+            magic_number: bytes
+            if isinstance(file.magic_number, str):
+                magic_number = bytes.fromhex(file.magic_number)
+            else:
+                magic_number = file.magic_number
             offset = getattr(file, "magic_number_offset", 0)
-            contents = os.urandom(offset)
-            magic_number = getattr(file, "magic_number", b"")
-            if isinstance(magic_number, str):
-                magic_number = bytes.fromhex(magic_number)
-            contents += magic_number
+            if offset < 0:
+                postamble = os.urandom(-(len(magic_number) + offset))
+                contents += magic_number + postamble  # type: ignore[operator]
+            else:
+                preamble = generator.generate_contents(binary=True, fill=offset)
+                contents = preamble + magic_number + contents  # type: ignore[operator]
         elif hasattr(file, "magic_pattern"):
             raise NotImplementedError(
                 "Sampling of magic version file types is not implemented yet"