Add a mechanism to register dataset builder providers.

This will allow to register dataset builders that are defined in different source code folders.

It also allows registering datasets that are defined elsewhere, e.g., on disk or in a database.

PiperOrigin-RevId: 669244349

Author: tomvdw

tensorflow_datasets/core/__init__.py

@@ -33,6 +33,8 @@
 from tensorflow_datasets.core.lazy_imports_lib import lazy_imports
 from tensorflow_datasets.core.load import DatasetCollectionLoader
 from tensorflow_datasets.core.naming import ShardedFileTemplate
+from tensorflow_datasets.core.registered import add_dataset_builder_provider
+from tensorflow_datasets.core.registered import DatasetBuilderProvider
 from tensorflow_datasets.core.registered import DatasetNotFoundError
 from tensorflow_datasets.core.sequential_writer import SequentialWriter
 from tensorflow_datasets.core.split_builder import SplitGeneratorLegacy as SplitGenerator
@@ -61,12 +63,14 @@ def benchmark(*args, **kwargs):
 
 __all__ = [
     "add_data_dir",
+    "add_dataset_builder_provider",
     "as_path",
     "BenchmarkResult",
     "BeamBasedBuilder",
     "BeamMetadataDict",
     "BuilderConfig",
     "DatasetBuilder",
+    "DatasetBuilderProvider",
     "DatasetCollectionLoader",
     "DatasetInfo",
     "DatasetIdentity",

tensorflow_datasets/core/registered.py

@@ -15,6 +15,8 @@
 
 """Access registered datasets."""
 
+from __future__ import annotations
+
 import abc
 import collections
 from collections.abc import Iterator
@@ -24,7 +26,7 @@
 import inspect
 import os.path
 import time
-from typing import ClassVar, Type
+from typing import ClassVar, Protocol, Type
 
 from absl import logging
 from etils import epath
@@ -37,11 +39,11 @@
 from tensorflow_datasets.core.utils import resource_utils
 
 # Internal registry containing <str registered_name, DatasetBuilder subclass>
-_DATASET_REGISTRY = {}
+_DATASET_REGISTRY: dict[str, Type[RegisteredDataset]] = {}
 
 # Internal registry containing:
 # <str snake_cased_name, abstract DatasetBuilder subclass>
-_ABSTRACT_DATASET_REGISTRY = {}
+_ABSTRACT_DATASET_REGISTRY: dict[str, Type[RegisteredDataset]] = {}
 
 # Keep track of dict[str (module name), list[DatasetBuilder]]
 # This is directly accessed by `tfds.community.builder_cls_from_module` when
@@ -289,6 +291,75 @@ def __init_subclass__(cls, skip_registration=False, **kwargs):  # pylint: disabl
       _DATASET_REGISTRY[cls.name] = cls
 
 
+class DatasetBuilderProvider(Protocol):
+
+  def has_dataset(self, name: str) -> bool:
+    ...
+
+  def get_builder_cls(self, name: str) -> Type[RegisteredDataset]:
+    ...
+
+
+class SourceDirDatasetBuilderProvider(DatasetBuilderProvider):
+  """Provider of dataset builders that are defined in the given source code folder."""
+
+  def __init__(self, datasets_dir: str):
+    self._datasets_dir = datasets_dir
+    self._registry: dict[str, Type[RegisteredDataset]] = {}
+
+  @functools.cached_property
+  def dataset_packages(self) -> dict[str, tuple[epath.Path, str]]:
+    """Returns existing datasets.
+
+    Returns:
+      {ds_name: (pkg_path, builder_module)}.
+      For example: {'mnist': ('/lib/tensorflow_datasets/datasets/mnist',
+                              'tensorflow_datasets.datasets.mnist.builder')}
+    """
+    datasets = {}
+    datasets_dir_path = resource_utils.tfds_path(self._datasets_dir)
+    if not datasets_dir_path.exists():
+      return datasets
+    ds_dir_pkg = '.'.join(
+        ['tensorflow_datasets'] + self._datasets_dir.split(os.path.sep)
+    )
+    for child in datasets_dir_path.iterdir():
+      # Except for a few exceptions, all children of datasets/ directory are
+      # packages of datasets, no needs to check child is a directory and
+      # contains a `builder.py` module.
+      exceptions = [
+          '__init__.py',
+      ]
+      if child.name not in exceptions:
+        pkg_path = epath.Path(datasets_dir_path) / child.name
+        builder_module = (
+            f'{ds_dir_pkg}.{child.name}.{child.name}{_BUILDER_MODULE_SUFFIX}'
+        )
+        datasets[child.name] = (pkg_path, builder_module)
+    return datasets
+
+  def has_dataset(self, name: str) -> bool:
+    return name in self.dataset_packages
+
+  def get_builder_cls(self, name: str) -> Type[RegisteredDataset]:
+    if name in self._registry:
+      return self._registry[name]
+    pkg_dir_path, builder_module = self.dataset_packages[name]
+    cls = importlib.import_module(builder_module).Builder
+    cls.pkg_dir_path = pkg_dir_path
+    self._registry[name] = cls
+    return cls
+
+
+_DATASET_PROVIDER_REGISTRY: list[DatasetBuilderProvider] = [
+    SourceDirDatasetBuilderProvider(constants.DATASETS_TFDS_SRC_DIR)
+]
+
+
+def add_dataset_builder_provider(provider: DatasetBuilderProvider) -> None:
+  _DATASET_PROVIDER_REGISTRY.append(provider)
+
+
 def _is_builder_available(builder_cls: Type[RegisteredDataset]) -> bool:
   """Returns `True` is the builder is available."""
   return visibility.DatasetType.TFDS_PUBLIC.is_available()
@@ -346,14 +417,9 @@ def _get_existing_dataset_packages(
 
 def imported_builder_cls(name: str) -> Type[RegisteredDataset]:
   """Returns the Registered dataset class."""
-  existing_ds_pkgs = _get_existing_dataset_packages(
-      constants.DATASETS_TFDS_SRC_DIR
-  )
-  if name in existing_ds_pkgs:
-    pkg_dir_path, builder_module = existing_ds_pkgs[name]
-    cls = importlib.import_module(builder_module).Builder
-    cls.pkg_dir_path = pkg_dir_path
-    return cls
+  for dataset_builder_provider in _DATASET_PROVIDER_REGISTRY:
+    if dataset_builder_provider.has_dataset(name):
+      return dataset_builder_provider.get_builder_cls(name)
 
   if name in _ABSTRACT_DATASET_REGISTRY:
     # Will raise TypeError: Can't instantiate abstract class X with abstract

tensorflow_datasets/core/registered_test.py

@@ -20,7 +20,6 @@
 from unittest import mock
 import pytest
 from tensorflow_datasets import testing
-from tensorflow_datasets.core import constants
 from tensorflow_datasets.core import load
 from tensorflow_datasets.core import registered
 from tensorflow_datasets.core import splits
@@ -326,14 +325,16 @@ def test_name_inferred_from_pkg_level3():
   assert ds_builder.name == "dummy_ds_2"
 
 
-class ConfigBasedBuildersTest(testing.TestCase):
+class SourceDirDatasetBuilderProviderTest(testing.TestCase):
 
-  def test__get_existing_dataset_packages(self):
-    ds_packages = registered._get_existing_dataset_packages(
+  def test_provider(self):
+    provider = registered.SourceDirDatasetBuilderProvider(
         "testing/dummy_config_based_datasets"
     )
-    self.assertEqual(set(ds_packages.keys()), {"dummy_ds_1", "dummy_ds_2"})
-    pkg_path, builder_module = ds_packages["dummy_ds_1"]
+    self.assertEqual(
+        set(provider.dataset_packages), {"dummy_ds_1", "dummy_ds_2"}
+    )
+    pkg_path, builder_module = provider.dataset_packages["dummy_ds_1"]
     self.assertEndsWith(
         str(pkg_path),
         "tensorflow_datasets/testing/dummy_config_based_datasets/dummy_ds_1",
@@ -343,10 +344,12 @@ def test__get_existing_dataset_packages(self):
         "tensorflow_datasets.testing.dummy_config_based_datasets.dummy_ds_1.dummy_ds_1_dataset_builder",
     )
 
-  @mock.patch.object(
-      constants, "DATASETS_TFDS_SRC_DIR", "testing/dummy_config_based_datasets"
-  )
   def test_imported_builder_cls(self):
+    registered.add_dataset_builder_provider(
+        registered.SourceDirDatasetBuilderProvider(
+            "testing/dummy_config_based_datasets"
+        )
+    )
     builder = registered.imported_builder_cls("dummy_ds_1")
     self.assertEqual(builder.name, "dummy_ds_1")