Merge pull request #55 from ArcanaFramework/sample-file-generator

refactored generate_sample_data

Author: tclose

extras/fileformats/extras/application/medical.py

@@ -1,10 +1,10 @@
 import typing as ty
 from pathlib import Path
-from random import Random
 import pydicom
 from fileformats.core import FileSet
 from fileformats.application import Dicom
 import medimages4tests.dummy.dicom.mri.t1w.siemens.skyra.syngo_d13c
+from fileformats.core.utils import SampleFileGenerator
 
 
 @FileSet.read_metadata.register
@@ -24,10 +24,10 @@ def dicom_read_metadata(
 @FileSet.generate_sample_data.register
 def dicom_generate_sample_data(
     dicom: Dicom,
-    dest_dir: Path,
-    seed: ty.Union[int, Random] = 0,
-    stem: ty.Optional[str] = None,
+    generator: SampleFileGenerator,
 ) -> ty.Iterable[Path]:
     return next(
-        medimages4tests.dummy.dicom.mri.t1w.siemens.skyra.syngo_d13c.get_image().iterdir()
+        medimages4tests.dummy.dicom.mri.t1w.siemens.skyra.syngo_d13c.get_image(
+            out_dir=generator.dest_dir
+        ).iterdir()
     )

fileformats/application/serialization.py

@@ -1,12 +1,11 @@
 import json
 import typing as ty
-from random import Random
 from pathlib import Path
 from fileformats.core import hook, DataType, FileSet
 from fileformats.core.mixin import WithClassifiers
 from ..generic import File
 from fileformats.core.exceptions import FormatMismatchError
-from fileformats.core.utils import gen_filename
+from fileformats.core.utils import SampleFileGenerator
 
 
 class Schema(DataType):
@@ -96,15 +95,17 @@ class Toml(DataSerialization):
 @FileSet.generate_sample_data.register
 def generate_json_sample_data(
     js: Json,
-    dest_dir: Path,
-    seed: ty.Union[int, Random] = 0,
-    stem: ty.Optional[str] = None,
+    generator: SampleFileGenerator,
 ) -> ty.Iterable[Path]:
-    js_file = dest_dir / gen_filename(seed, file_type=js, stem=stem)
-    rng = Random(seed + 1)
+    js_file = generator.generate_fspath(file_type=Json)
     with open(js_file, "w") as f:
         json.dump(
-            {"a": True, "b": "two", "c": 3, "d": [rng.randint(0, 10), rng.random(), 6]},
+            {
+                "a": True,
+                "b": "two",
+                "c": 3,
+                "d": [generator.rng.randint(0, 10), generator.rng.random(), 6],
+            },
             f,
         )
     return [js_file]
@@ -113,21 +114,18 @@ def generate_json_sample_data(
 @FileSet.generate_sample_data.register
 def generate_yaml_sample_data(
     yml: Yaml,
-    dest_dir: Path,
-    seed: ty.Union[int, Random] = 0,
-    stem: ty.Optional[str] = None,
+    generator: SampleFileGenerator,
 ) -> ty.Iterable[Path]:
-    yml_file = dest_dir / gen_filename(seed, file_type=yml, stem=stem)
-    rng = Random(seed + 1)
+    yml_file = generator.generate_fspath(file_type=Yaml)
     with open(yml_file, "w") as f:
         f.write(
             f"""# Generated sample YAML file by FileFormats
 a: True
 b: two
 c: 3
 d:
-- {rng.randint(0, 10)}
-- {rng.random()}
+- {generator.rng.randint(0, 10)}
+- {generator.rng.random()}
 - 6
 """
         )

fileformats/core/fileset.py

@@ -4,7 +4,6 @@
 from enum import Enum, IntEnum
 from warnings import warn
 import tempfile
-from random import Random
 from collections import Counter
 import typing as ty
 import shutil
@@ -23,6 +22,7 @@
     describe_task,
     matching_source,
     import_extras_module,
+    SampleFileGenerator,
 )
 from .converter import SubtypeVar
 from .classifier import Classifier
@@ -859,9 +859,9 @@ def referenced_types(cls) -> ty.Set[Classifier]:
         return types
 
     @classmethod
-    def mock(cls, *fspaths: ty.Tuple[ty.Union[Path, str]]) -> Self:
+    def mock(cls, *fspaths: ty.Tuple[ty.Union[Path, str]]) -> "FileSet":
         """Return an instance of a mocked sub-class of the file format to be used in
-        test routines like doctests.
+        test routines like doctests that doesn't require to point at actual files
 
         Parameters
         ----------
@@ -909,8 +909,9 @@ def sample(
             dest_dir = Path(tempfile.mkdtemp())
         # Need to use mock to get an instance in order to use the singledispatch-based
         # hook.extra decorator
-        mock = cls.mock()
-        fspaths = mock.generate_sample_data(dest_dir, seed, stem)
+        fspaths = cls.sample_data(
+            SampleFileGenerator(dest_dir=dest_dir, seed=seed, fname_stem=stem)
+        )
         try:
             obj = cls(fspaths)
         except FormatMismatchError as e:
@@ -921,12 +922,28 @@ def sample(
             )
         return obj
 
+    @classmethod
+    def sample_data(cls, generator: SampleFileGenerator) -> ty.Iterable[Path]:
+        """Converts the `generate_sample_data` method into a class method by mocking up
+        a class instance and calling the method on it
+
+        Parameters
+        ----------
+        generator : SampleFileGenerator
+            the generator to use to create the sample data
+
+        Returns
+        -------
+        ty.Iterable[Path]
+            the generated file-system paths
+        """
+        mock: FileSet = cls.mock()
+        return mock.generate_sample_data(generator)
+
     @hook.extra
     def generate_sample_data(
         self,
-        dest_dir: Path,
-        seed: ty.Union[int, Random] = 0,
-        stem: ty.Optional[str] = None,
+        generator: SampleFileGenerator,
     ) -> ty.Iterable[Path]:
         """Generate test data at the fspaths of the file-set
 

fileformats/core/utils.py

@@ -5,6 +5,7 @@
 import random
 import string
 import inspect
+from functools import cached_property
 import typing as ty
 import re
 import urllib.request
@@ -488,45 +489,157 @@ def import_extras_module(klass: type) -> ExtrasModule:
 LIST_MIME = "+list-of"
 
 
-def gen_filename(
-    seed_or_rng: ty.Union[random.Random, int],
-    file_type: ty.Type[fileformats.core.FileSet] = None,
-    length: int = 32,
-    stem: ty.Optional[str] = None,
-):
-    """Generates a random filename of length `length` and extension `ext`
+class SampleFileGenerator:
+    """Generates sample files. Designed to be used within generate_sample_data overrides
 
     Parameters
     ----------
-    seed_or_rng : random.Random or int
-        used to seed the random number generator
-    file_type : Type[FileSet], optional
-        type of the file to generate the filename for, used to append any extensions
-        and seed the random number generator if required
-    length : int
-        length of the filename (minus extension)
-    stem : str, optional
-        the stem to use for the filename if provided
-
-    Returns
-    -------
-    filename : str
-        randomly generated filename
+    dest_dir : Path
+        the directory to write the sample files to
+    seed : int
+        the seed for the random number generator
+    fname_stem : str
+        the stem of the file name to generate
     """
-    if file_type is None:
-        import fileformats.generic
 
-        file_type = fileformats.generic.FsObject
-    if stem:
-        fname = stem
-    else:
-        if isinstance(seed_or_rng, random.Random):
-            rng = seed_or_rng
+    dest_dir: Path
+    seed: int
+    fname_stem: str
+
+    FNAME_STEM_LENGTH = 24
+
+    def __init__(self, dest_dir: Path, seed: int, fname_stem: str = None):
+        self.dest_dir = dest_dir
+        self.seed = seed
+        self.fname_stem = (
+            self._generate_fname_stem() if fname_stem is None else fname_stem
+        )
+
+    def _generate_fname_stem(self):
+        return "".join(
+            self.rng.choices(
+                string.ascii_letters + string.digits, k=self.FNAME_STEM_LENGTH
+            )
+        )
+
+    @cached_property
+    def rng(self):
+        return random.Random(self.seed)
+
+    def generate(
+        self,
+        file_type: ty.Type[fileformats.core.FileSet],
+        contents: ty.Union[str, bytes] = None,
+        fill: int = 0,
+        **kwargs,
+    ):
+        """Generates a random file of length `length` and extension `ext`
+
+        Parameters
+        ----------
+        file_type : Type[FileSet]
+            type of the file to generate the filename for, used to append any extensions
+            and seed the random number generator if required
+        contents : Union[str, bytes]
+            the contents of the file to write
+        fill : int
+            length of the random string to generate for the file contents. Will be appended
+            after any explicitly provided contents
+        **kwargs : dict
+            additional keyword arguments to pass to generate_fspath
+
+        Returns
+        -------
+        fspath : Path
+            path to the randomly generated file
+        """
+        if not contents and not fill:
+            raise ValueError("Either contents or random_fill_length must be provided")
+        fspath = self.generate_fspath(file_type, **kwargs)
+        fspath.parent.mkdir(parents=True, exist_ok=True)
+        try:
+            is_binary = file_type.binary
+        except AttributeError:
+            is_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))
+            )
+        else:
+            contents_type = bytes if is_binary else str
+            if not isinstance(contents, bytes):
+                raise TypeError(
+                    f"contents must be {contents_type} for {file_type} files, "
+                    f"not {type(contents)}"
+                )
+        if is_binary:
+            fspath.write_bytes(contents)
         else:
-            if not inspect.isclass(file_type):
-                file_type = type(file_type)
-            rng = random.Random(str(seed_or_rng) + file_type.mime_like)
-        fname = "".join(rng.choices(string.ascii_letters + string.digits, k=length))
-    if file_type and file_type.ext:
-        fname += file_type.ext
-    return fname
+            fspath.write_text(contents)
+        return fspath
+
+    def generate_fspath(
+        self,
+        file_type: ty.Optional[ty.Type[fileformats.core.FileSet]] = None,
+        fname_stem: ty.Optional[str] = None,
+        relpath: ty.Optional[Path] = None,
+    ):
+        """Generates a random file path in the destination directory of length `length`
+        and extension `ext`
+
+        Parameters
+        ----------
+        file_type : Type[FileSet]
+            type of the file to generate the filename for, used to append any extensions
+            and seed the random number generator if required
+        fname_stem : str, optional or bool
+            Use explicitly provided if it is a string
+        relpath : Path
+            the path to generate the filename at, relative to the destination directory
+
+        Returns
+        -------
+        fspath : Path
+            randomly generated file-system path
+        """
+        if file_type is None:
+            import fileformats.generic
+
+            file_type = fileformats.generic.FsObject
+        if fname_stem is not None:
+            fname = fname_stem
+        else:
+            fname = self.fname_stem
+        if file_type and file_type.ext:
+            fname += file_type.ext
+        fspath = self.dest_dir
+        if relpath:
+            fspath /= relpath
+        return fspath / fname
+
+    def child(
+        self, dest_dir: ty.Optional[Path] = None, fname_stem: str = None
+    ) -> "SampleFileGenerator":
+        """Creates a new instance of SampleFileGenerator with the same destination
+        directory and seed, but a new random filename stem
+
+        Parameters
+        ----------
+        relpath : Path, optional
+            the path to generate the filename at, relative to the destination directory
+        fname_stem : str, optional
+            the stem of the file name to generate
+
+        Returns
+        -------
+        SampleFileGenerator
+            the new instance of SampleFileGenerator
+        """
+        if dest_dir is None:
+            dest_dir = self.dest_dir
+        kwargs = {"fname_stem": fname_stem} if fname_stem else {}
+        return SampleFileGenerator(
+            dest_dir, seed=self.rng.randint(0, 2**32 - 1), **kwargs
+        )