[PoC] merge mock data preparation and loading

Author: pmeier

test/builtin_dataset_mocks.py

@@ -10,6 +10,7 @@
 import pathlib
 import pickle
 import random
+import shutil
 import unittest.mock
 import warnings
 import xml.etree.ElementTree as ET
@@ -22,7 +23,6 @@
 from datasets_utils import make_zip, make_tar, create_image_folder, create_image_file, combinations_grid
 from torch.nn.functional import one_hot
 from torch.testing import make_tensor as _make_tensor
-from torchvision._utils import sequence_to_str
 from torchvision.prototype import datasets
 
 make_tensor = functools.partial(_make_tensor, device="cpu")
@@ -62,27 +62,47 @@ def _parse_mock_info(self, mock_info):
 
         return mock_info
 
-    def prepare(self, config):
+    def load(self, config):
         # `datasets.home()` is patched to a temporary directory through the autouse fixture `test_home` in
         # test/test_prototype_builtin_datasets.py
         root = pathlib.Path(datasets.home()) / self.name
-        root.mkdir(exist_ok=True)
+        mock_data_folder = root / "__mock__"
+        mock_data_folder.mkdir(parents=True)
 
-        mock_info = self._parse_mock_info(self.mock_data_fn(root, config))
+        mock_info = self._parse_mock_info(self.mock_data_fn(mock_data_folder, config))
 
-        with unittest.mock.patch.object(datasets.utils.Dataset, "__init__"):
-            required_file_names = {
-                resource.file_name for resource in datasets.load(self.name, root=root, **config)._resources()
-            }
-        available_file_names = {path.name for path in root.glob("*")}
-        missing_file_names = required_file_names - available_file_names
-        if missing_file_names:
-            raise pytest.UsageError(
-                f"Dataset '{self.name}' requires the files {sequence_to_str(sorted(missing_file_names))} "
-                f"for {config}, but they were not created by the mock data function."
-            )
+        def mock_data_download(resource, root, **kwargs):
+            src = mock_data_folder / resource.file_name
+            if not src.exists():
+                raise pytest.UsageError(
+                    f"Dataset '{self.name}' requires the file {resource.file_name} for {config}"
+                    f"but it was not created by the mock data function."
+                )
 
-        return mock_info
+            dst = root / resource.file_name
+            shutil.move(str(src), str(root))
+
+            return dst
+
+        with unittest.mock.patch(
+            "torchvision.prototype.datasets.utils._resource.OnlineResource.download", new=mock_data_download
+        ):
+            dataset = datasets.load(self.name, **config)
+
+        extra_files = list(mock_data_folder.glob("**/*"))
+        if not extra_files:
+            mock_data_folder.rmdir()
+        else:
+            pass
+            # raise pytest.UsageError(
+            #     (
+            #         f"Dataset '{self.name}' created the following files for {config} in the mock data function, "
+            #         f"but they were not loaded:\n\n"
+            #     )
+            #     + "\n".join(str(file.relative_to(mock_data_folder)) for file in extra_files)
+            # )
+
+        return dataset, mock_info
 
 
 def config_id(name, config):

test/test_prototype_builtin_datasets.py

@@ -56,18 +56,14 @@ def test_info(self, name):
 
     @parametrize_dataset_mocks(DATASET_MOCKS)
     def test_smoke(self, dataset_mock, config):
-        dataset_mock.prepare(config)
-
-        dataset = datasets.load(dataset_mock.name, **config)
+        dataset, _ = dataset_mock.load(config)
 
         if not isinstance(dataset, datasets.utils.Dataset):
             raise AssertionError(f"Loading the dataset should return an Dataset, but got {type(dataset)} instead.")
 
     @parametrize_dataset_mocks(DATASET_MOCKS)
     def test_sample(self, dataset_mock, config):
-        dataset_mock.prepare(config)
-
-        dataset = datasets.load(dataset_mock.name, **config)
+        dataset, _ = dataset_mock.load(config)
 
         try:
             sample = next(iter(dataset))
@@ -84,17 +80,13 @@ def test_sample(self, dataset_mock, config):
 
     @parametrize_dataset_mocks(DATASET_MOCKS)
     def test_num_samples(self, dataset_mock, config):
-        mock_info = dataset_mock.prepare(config)
-
-        dataset = datasets.load(dataset_mock.name, **config)
+        dataset, mock_info = dataset_mock.load(config)
 
         assert len(list(dataset)) == mock_info["num_samples"]
 
     @parametrize_dataset_mocks(DATASET_MOCKS)
     def test_no_vanilla_tensors(self, dataset_mock, config):
-        dataset_mock.prepare(config)
-
-        dataset = datasets.load(dataset_mock.name, **config)
+        dataset, _ = dataset_mock.load(config)
 
         vanilla_tensors = {key for key, value in next(iter(dataset)).items() if type(value) is torch.Tensor}
         if vanilla_tensors:
@@ -105,24 +97,20 @@ def test_no_vanilla_tensors(self, dataset_mock, config):
 
     @parametrize_dataset_mocks(DATASET_MOCKS)
     def test_transformable(self, dataset_mock, config):
-        dataset_mock.prepare(config)
-
-        dataset = datasets.load(dataset_mock.name, **config)
+        dataset, _ = dataset_mock.load(config)
 
         next(iter(dataset.map(transforms.Identity())))
 
     @pytest.mark.parametrize("only_datapipe", [False, True])
     @parametrize_dataset_mocks(DATASET_MOCKS)
     def test_traversable(self, dataset_mock, config, only_datapipe):
-        dataset_mock.prepare(config)
-        dataset = datasets.load(dataset_mock.name, **config)
+        dataset, _ = dataset_mock.load(config)
 
         traverse(dataset, only_datapipe=only_datapipe)
 
     @parametrize_dataset_mocks(DATASET_MOCKS)
     def test_serializable(self, dataset_mock, config):
-        dataset_mock.prepare(config)
-        dataset = datasets.load(dataset_mock.name, **config)
+        dataset, _ = dataset_mock.load(config)
 
         pickle.dumps(dataset)
 
@@ -135,8 +123,7 @@ def _collate_fn(self, batch):
     @pytest.mark.parametrize("num_workers", [0, 1])
     @parametrize_dataset_mocks(DATASET_MOCKS)
     def test_data_loader(self, dataset_mock, config, num_workers):
-        dataset_mock.prepare(config)
-        dataset = datasets.load(dataset_mock.name, **config)
+        dataset, _ = dataset_mock.load(config)
 
         dl = DataLoader(
             dataset,
@@ -153,17 +140,15 @@ def test_data_loader(self, dataset_mock, config, num_workers):
     @parametrize_dataset_mocks(DATASET_MOCKS)
     @pytest.mark.parametrize("annotation_dp_type", (Shuffler, ShardingFilter))
     def test_has_annotations(self, dataset_mock, config, annotation_dp_type):
-
-        dataset_mock.prepare(config)
-        dataset = datasets.load(dataset_mock.name, **config)
+        dataset, _ = dataset_mock.load(config)
 
         if not any(isinstance(dp, annotation_dp_type) for dp in extract_datapipes(dataset)):
             raise AssertionError(f"The dataset doesn't contain a {annotation_dp_type.__name__}() datapipe.")
 
     @parametrize_dataset_mocks(DATASET_MOCKS)
     def test_save_load(self, dataset_mock, config):
-        dataset_mock.prepare(config)
-        dataset = datasets.load(dataset_mock.name, **config)
+        dataset, _ = dataset_mock.load(config)
+
         sample = next(iter(dataset))
 
         with io.BytesIO() as buffer:
@@ -173,8 +158,7 @@ def test_save_load(self, dataset_mock, config):
 
     @parametrize_dataset_mocks(DATASET_MOCKS)
     def test_infinite_buffer_size(self, dataset_mock, config):
-        dataset_mock.prepare(config)
-        dataset = datasets.load(dataset_mock.name, **config)
+        dataset, _ = dataset_mock.load(config)
 
         for dp in extract_datapipes(dataset):
             if hasattr(dp, "buffer_size"):
@@ -184,18 +168,15 @@ def test_infinite_buffer_size(self, dataset_mock, config):
 
     @parametrize_dataset_mocks(DATASET_MOCKS)
     def test_has_length(self, dataset_mock, config):
-        dataset_mock.prepare(config)
-        dataset = datasets.load(dataset_mock.name, **config)
+        dataset, _ = dataset_mock.load(config)
 
         assert len(dataset) > 0
 
 
 @parametrize_dataset_mocks(DATASET_MOCKS["qmnist"])
 class TestQMNIST:
     def test_extra_label(self, dataset_mock, config):
-        dataset_mock.prepare(config)
-
-        dataset = datasets.load(dataset_mock.name, **config)
+        dataset, _ = dataset_mock.load(config)
 
         sample = next(iter(dataset))
         for key, type in (
@@ -218,9 +199,7 @@ def test_label_matches_path(self, dataset_mock, config):
         if config["split"] != "train":
             return
 
-        dataset_mock.prepare(config)
-
-        dataset = datasets.load(dataset_mock.name, **config)
+        dataset, _ = dataset_mock.load(config)
 
         for sample in dataset:
             label_from_path = int(Path(sample["path"]).parent.name)
@@ -230,9 +209,7 @@ def test_label_matches_path(self, dataset_mock, config):
 @parametrize_dataset_mocks(DATASET_MOCKS["usps"])
 class TestUSPS:
     def test_sample_content(self, dataset_mock, config):
-        dataset_mock.prepare(config)
-
-        dataset = datasets.load(dataset_mock.name, **config)
+        dataset, _ = dataset_mock.load(config)
 
         for sample in dataset:
             assert "image" in sample