Add builders defined in packages only when TFDS PUBLIC is visible

PiperOrigin-RevId: 713279260

Author: tomvdw

tensorflow_datasets/core/registered.py

@@ -422,7 +422,11 @@ def list_imported_builders() -> list[str]:
       builder_name
       for builder_name, builder_cls in _DATASET_REGISTRY.items()
       if _is_builder_available(builder_cls)
-  ] + list(_get_existing_dataset_packages(constants.DATASETS_TFDS_SRC_DIR))
+  ]
+  if visibility.DatasetType.TFDS_PUBLIC.is_available():
+    all_builders += list(
+        _get_existing_dataset_packages(constants.DATASETS_TFDS_SRC_DIR)
+    )
   return sorted(all_builders)
 
 

tensorflow_datasets/core/visibility_test.py

@@ -13,8 +13,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Tests for tensorflow_datasets.core.visibility."""
-
 import tensorflow_datasets as tfds
 from tensorflow_datasets.core import visibility
 
@@ -28,3 +26,9 @@ def test_visibility():
   # `absl.app` should detect the TFDS script and restrict the visibility
   # to TFDS public by default.
   assert visibility._current_available == {visibility.DatasetType.TFDS_PUBLIC}
+  assert 'mnist' in tfds.list_builders()
+  assert 'abstract_reasoning' in tfds.list_builders()
+  # Remove all visibility and check that the datasets are not available.
+  with visibility.set_availables_tmp([]):
+    assert 'mnist' not in tfds.list_builders()
+    assert 'abstract_reasoning' not in tfds.list_builders()