Add a functionality to reset the list of dataset builder providers.BEGIN PUBLIC

PiperOrigin-RevId: 684549116

Author: The TensorFlow Datasets Authors

tensorflow_datasets/core/registered.py

@@ -371,10 +371,16 @@ def get_builder_cls(self, name: str) -> Type[RegisteredDataset]:
     return cls
 
 
-_DATASET_PROVIDER_REGISTRY: list[DatasetBuilderProvider] = [
-    SourceDirDatasetBuilderProvider(constants.DATASETS_TFDS_SRC_DIR),
-    LegacyDatasetBuilderProvider(),
-]
+def _get_inital_providers() -> list[DatasetBuilderProvider]:
+  return [
+      SourceDirDatasetBuilderProvider(constants.DATASETS_TFDS_SRC_DIR),
+      LegacyDatasetBuilderProvider(),
+  ]
+
+
+_DATASET_PROVIDER_REGISTRY: list[DatasetBuilderProvider] = (
+    _get_inital_providers()
+)
 
 
 def add_dataset_builder_provider(
@@ -394,6 +400,12 @@ def add_dataset_builder_provider(
     _DATASET_PROVIDER_REGISTRY.append(provider)
 
 
+def reset_dataset_builder_providers() -> None:
+  """Resets the list of dataset builder providers to remove added providers."""
+  global _DATASET_PROVIDER_REGISTRY
+  _DATASET_PROVIDER_REGISTRY = _get_inital_providers()
+
+
 def _is_builder_available(builder_cls: Type[RegisteredDataset]) -> bool:
   """Returns `True` is the builder is available."""
   return visibility.DatasetType.TFDS_PUBLIC.is_available()

tensorflow_datasets/core/registered_test.py

@@ -287,6 +287,28 @@ class MyDataset(registered.RegisteredDataset):  # pylint: disable=unused-variabl
     builder_cls = registered.imported_builder_cls("my_dataset")
     self.assertEqual(builder_cls, EmptyDatasetBuilder)
 
+  def test_reset_dataset_builder_providers(self):
+    """Resetting dataset builder providers."""
+
+    class MyRegisteredDataset(registered.RegisteredDataset):  # pylint: disable=unused-variable
+      pass
+
+    registered.add_dataset_builder_provider(
+        TestBuilderProvider("my_registered_dataset", EmptyDatasetBuilder),
+        0,  # start of the list
+    )
+    builder_cls = registered.imported_builder_cls("my_registered_dataset")
+    self.assertEqual(builder_cls, EmptyDatasetBuilder)
+
+    registered.reset_dataset_builder_providers()
+
+    registered.add_dataset_builder_provider(
+        TestBuilderProvider("my_registered_dataset", EmptyDatasetBuilder),
+        None,  # end of the list
+    )
+    builder_cls = registered.imported_builder_cls("my_registered_dataset")
+    self.assertEqual(builder_cls, MyRegisteredDataset)
+
 
 def test_skip_regitration():
   """Test `skip_registration()`."""