more robust way of inferring Builder name, support builder pkg located anywhere.

PiperOrigin-RevId: 502523472

Author: pierrot0

tensorflow_datasets/core/registered.py

@@ -54,6 +54,9 @@
 # Keep track of Dict[str (module name), List[DatasetCollectionBuilder]]
 _MODULE_TO_DATASET_COLLECTIONS = collections.defaultdict(list)
 
+# eg for dataset "foo": "tensorflow_datasets.datasets.foo.foo_dataset_builder".
+_BUILDER_MODULE_SUFFIX = '_dataset_builder'
+
 
 class DatasetNotFoundError(ValueError):
   """Exception raised when the dataset cannot be found."""
@@ -169,12 +172,19 @@ def __init_subclass__(cls, skip_registration=False, **kwargs):  # pylint: disabl
     if not cls.__dict__.get('name'):
       if cls.__name__ == 'Builder':
         # Config-based builders should be defined with a class named "Builder".
-        # In such a case, the builder name is extracted from the package name:
-        # 1 package = 1 dataset!.
-        module_path_components = cls.__module__.rsplit('.', 2)
-        if len(module_path_components) != 3:
-          raise AssertionError(f'module path is too short: {cls.__module__}')
-        cls.name = module_path_components[1]
+        # In such a case, the builder name is extracted from the module if it
+        # follows conventions:
+        module_name = cls.__module__.rsplit('.', 1)[-1]
+        if module_name.endswith(_BUILDER_MODULE_SUFFIX):
+          cls.name = module_name[: -len(_BUILDER_MODULE_SUFFIX)]
+        elif '.' in cls.__module__:  # Extract dataset name from package name.
+          cls.name = cls.__module__.rsplit('.', 2)[-2]
+        else:
+          raise AssertionError(
+              'When using `Builder` as class name, the dataset builder name is '
+              'inferred from module name if named "*_dataset_builder" or from '
+              f'package name, but there is no package in "{cls.__module__}".'
+          )
       else:  # Legacy builders.
         cls.name = naming.camelcase_to_snakecase(cls.__name__)
 
@@ -261,7 +271,9 @@ def _get_existing_dataset_packages(
     ]
     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}_dataset_builder'
+      builder_module = (
+          f'{ds_dir_pkg}.{child.name}.{child.name}{_BUILDER_MODULE_SUFFIX}'
+      )
       datasets[child.name] = (pkg_path, builder_module)
   return datasets
 

tensorflow_datasets/core/registered_test.py

@@ -16,6 +16,8 @@
 """Tests for tensorflow_datasets.core.registered."""
 
 import abc
+import re
+
 from unittest import mock
 import pytest
 
@@ -26,6 +28,7 @@
 from tensorflow_datasets.core import splits
 from tensorflow_datasets.core import utils
 from tensorflow_datasets.core.utils import py_utils
+import tensorflow_datasets.public_api as tfds
 from tensorflow_datasets.testing.dummy_config_based_datasets.dummy_ds_1 import dummy_ds_1_dataset_builder
 
 
@@ -278,13 +281,57 @@ def test_name_inferred_from_pkg():
   assert ds_builder.name == "dummy_ds_1"
 
 
+def test_name_inferred_from_pkg_level0_fails():
+  pkg_path = (
+      tfds.core.tfds_path() / "testing/dummy_config_based_datasets/dummy_ds_2"
+  )
+  expected_msg = (
+      "When using `Builder` as class name, the dataset builder name is "
+      'inferred from module name if named "*_dataset_builder" or from '
+      'package name, but there is no package in "dummy_builder".'
+  )
+  with tfds.core.utils.add_sys_path(pkg_path):
+    with pytest.raises(AssertionError, match=re.escape(expected_msg)):
+      tfds.core.community.builder_cls_from_module("dummy_builder")
+
+
+@mock.patch.dict(registered._DATASET_REGISTRY, {})
+def test_name_inferred_from_pkg_level1():
+  pkg_path = tfds.core.tfds_path() / "testing/dummy_config_based_datasets"
+  with tfds.core.utils.add_sys_path(pkg_path):
+    ds_builder = tfds.core.community.builder_cls_from_module(
+        "dummy_ds_2.dummy_builder"
+    )
+  assert ds_builder.name == "dummy_ds_2"
+
+
+@mock.patch.dict(registered._DATASET_REGISTRY, {})
+def test_name_inferred_from_pkg_level2():
+  pkg_path = tfds.core.tfds_path() / "testing"
+  with tfds.core.utils.add_sys_path(pkg_path):
+    ds_builder = tfds.core.community.builder_cls_from_module(
+        "dummy_config_based_datasets.dummy_ds_2.dummy_builder"
+    )
+  assert ds_builder.name == "dummy_ds_2"
+
+
+@mock.patch.dict(registered._DATASET_REGISTRY, {})
+def test_name_inferred_from_pkg_level3():
+  pkg_path = tfds.core.tfds_path()
+  with tfds.core.utils.add_sys_path(pkg_path):
+    ds_builder = tfds.core.community.builder_cls_from_module(
+        "testing.dummy_config_based_datasets.dummy_ds_2.dummy_builder"
+    )
+  assert ds_builder.name == "dummy_ds_2"
+
+
 class ConfigBasedBuildersTest(testing.TestCase):
 
   def test__get_existing_dataset_packages(self):
     ds_packages = registered._get_existing_dataset_packages(
         "testing/dummy_config_based_datasets"
     )
-    self.assertEqual(list(ds_packages.keys()), ["dummy_ds_1"])
+    self.assertEqual(list(ds_packages.keys()), ["dummy_ds_1", "dummy_ds_2"])
     pkg_path, builder_module = ds_packages["dummy_ds_1"]
     self.assertEndsWith(
         str(pkg_path),

tensorflow_datasets/testing/dummy_config_based_datasets/dummy_ds_2/README.md

@@ -0,0 +1,2 @@
+Similar to `dummy_ds_1`, but the builder module does not follow
+the recommended naming.

tensorflow_datasets/testing/dummy_config_based_datasets/dummy_ds_2/dummy_builder.py

@@ -0,0 +1,48 @@
+# coding=utf-8
+# Copyright 2022 The TensorFlow Datasets Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Dummy config-based dataset self-contained in a directory.
+
+The builder module intentionnaly does not follow the naming conventions.
+This is used in registered_test.py to check the logic to infer dataset
+name.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import tensorflow_datasets.public_api as tfds
+
+
+class Builder(tfds.core.GeneratorBasedBuilder):
+  """Dummy dataset."""
+
+  VERSION = tfds.core.Version('1.0.0')
+
+  def _info(self):
+    return self.dataset_info_from_configs(
+        features=tfds.features.FeaturesDict({'x': np.int64}),
+    )
+
+  def _split_generators(self, dl_manager):
+    return [
+        tfds.core.SplitGenerator(
+            name=tfds.Split.TRAIN,
+        ),
+    ]
+
+  def _generate_examples(self):
+    for i in range(10):
+      yield i, {'x': i}