Fix lazy importing for dataset tests (#3481)

test/datasets_utils.py

@@ -10,6 +10,7 @@ import random
 import string
 import unittest
 import unittest.mock
+from collections import defaultdict
 from typing import Any, Callable, Dict, Iterator, List, Optional, Sequence, Tuple, Union
 
 import PIL
@@ -60,26 +61,42 @@ class LazyImporter:
     )
 
     def __init__(self):
-        cls = type(self)
+        modules = defaultdict(list)
         for module in self.MODULES:
-            # We need the quirky 'module=module' argument to the lambda since otherwise the lookup for 'module' in this
-            # scope would happen at runtime rather than at definition. Thus, without it, every property would try to
-            # import the last 'module' in MODULES.
-            setattr(cls, module.split(".", 1)[0], property(lambda self, module=module: LazyImporter._import(module)))
+            module, *submodules = module.split(".", 1)
+            if submodules:
+                modules[module].append(submodules[0])
+            else:
+                # This introduces the module so that it is known when we later iterate over the dictionary.
+                modules.__missing__(module)
+
+        for module, submodules in modules.items():
+            # We need the quirky 'module=module' and submodules=submodules arguments to the lambda since otherwise the
+            # lookup for these would happen at runtime rather than at definition. Thus, without it, every property
+            # would try to import the last item in 'modules'
+            setattr(
+                type(self),
+                module,
+                property(lambda self, module=module, submodules=submodules: LazyImporter._import(module, submodules)),
+            )
 
     @staticmethod
-    def _import(module):
+    def _import(package, subpackages):
         try:
-            importlib.import_module(module)
-            return importlib.import_module(module.split(".", 1)[0])
+            module = importlib.import_module(package)
         except ImportError as error:
             raise UsageError(
-                f"Failed to import module '{module}'. "
-                f"This probably means that the current test case needs '{module}' installed, "
+                f"Failed to import module '{package}'. "
+                f"This probably means that the current test case needs '{package}' installed, "
                 f"but it is not a dependency of torchvision. "
-                f"You need to install it manually, for example 'pip install {module}'."
+                f"You need to install it manually, for example 'pip install {package}'."
             ) from error
 
+        for name in subpackages:
+            importlib.import_module(f".{name}", package=package)
+
+        return module
+
 
 lazy_importer = LazyImporter()