Lazy load namespace config

PiperOrigin-RevId: 627964344

Author: tomvdw

tensorflow_datasets/core/community/config.py

@@ -17,6 +17,7 @@
 
 from collections.abc import Mapping, Sequence
 import dataclasses
+import functools
 import threading
 from typing import Any
 
@@ -114,6 +115,9 @@ def add_namespace(self, namespace: str, config: NamespaceConfig) -> None:
     self._config_per_namespace[namespace] = config
 
 
-community_config = NamespaceRegistry(
-    config_path=utils.tfds_path('community-datasets.toml')
-)
+COMMUNITY_CONFIG_PATH = 'community-datasets.toml'
+
+
+@functools.lru_cache(maxsize=1)
+def get_community_config() -> NamespaceRegistry:
+  return NamespaceRegistry(config_path=utils.tfds_path(COMMUNITY_CONFIG_PATH))

tensorflow_datasets/core/community/config_test.py

@@ -17,6 +17,7 @@
 
 
 def test_community_config():
-  assert "huggingface" in config_lib.community_config.config_per_namespace
-  assert "kubric" in config_lib.community_config.config_per_namespace
-  assert "robotics" in config_lib.community_config.config_per_namespace
+  community_config = config_lib.get_community_config()
+  assert "huggingface" in community_config.config_per_namespace
+  assert "kubric" in community_config.config_per_namespace
+  assert "robotics" in community_config.config_per_namespace

tensorflow_datasets/core/community/registry.py

@@ -18,6 +18,7 @@
 from collections.abc import Iterable, Mapping, Sequence
 import concurrent.futures
 import difflib
+import functools
 import multiprocessing
 import os
 from typing import Any, Type
@@ -308,4 +309,6 @@ def get_builder_root_dirs(
     return result
 
 
-community_register = DatasetRegistry(config_lib.community_config)
+@functools.lru_cache(maxsize=1)
+def community_register() -> DatasetRegistry:
+  return DatasetRegistry(config_lib.get_community_config())

tensorflow_datasets/core/community/registry_test.py

@@ -153,7 +153,7 @@ def test_load_register_for_path_mixed():
 
 
 def test_community_register():
-  assert 'huggingface' in registry_lib.community_register.list_namespaces()
+  assert 'huggingface' in registry_lib.community_register().list_namespaces()
 
 
 def _write_dummy_config(content: str, tmp_path: str) -> epath.Path:

tensorflow_datasets/core/load.py

@@ -78,7 +78,7 @@ def list_builders(
   datasets = registered.list_imported_builders()
   if with_community_datasets:
     if visibility.DatasetType.COMMUNITY_PUBLIC.is_available():
-      datasets += community.community_register.list_builders()
+      datasets += community.community_register().list_builders()
   return datasets
 
 
@@ -112,7 +112,7 @@ def builder_cls(name: str) -> Type[dataset_builder.DatasetBuilder]:
   if ds_name.namespace:
     # `namespace:dataset` are loaded from the community register
     if visibility.DatasetType.COMMUNITY_PUBLIC.is_available():
-      return community.community_register.builder_cls(ds_name)
+      return community.community_register().builder_cls(ds_name)
     else:
       raise ValueError(
           f'Cannot load {ds_name} when community datasets are disabled'
@@ -191,9 +191,9 @@ def get_dataset_repr() -> str:
           name=name.name, **builder_kwargs)
     if (
         visibility.DatasetType.COMMUNITY_PUBLIC.is_available()
-        and community.community_register.has_namespace(name.namespace)
+        and community.community_register().has_namespace(name.namespace)
     ):
-      return community.community_register.builder(name=name, **builder_kwargs)
+      return community.community_register().builder(name=name, **builder_kwargs)
 
   # First check whether we can find the corresponding dataset builder code
   try:

tensorflow_datasets/core/load_test.py

@@ -45,10 +45,10 @@ def test_load_hf_dataset():
 ])
 def test_community_public_load():
   with mock.patch(
-      'tensorflow_datasets.core.community.community_register.list_builders',
+      'tensorflow_datasets.core.community.registry.DatasetRegistry.list_builders',
       return_value=['ns:ds'],
   ), mock.patch(
-      'tensorflow_datasets.core.community.community_register.builder_cls',
+      'tensorflow_datasets.core.community.registry.DatasetRegistry.builder_cls',
       return_value=testing.DummyDataset,
   ), mock.patch(
       'tensorflow_datasets.core.registered.list_imported_builders',