Lazily load various libraries

PiperOrigin-RevId: 627961527

Author: tomvdw

tensorflow_datasets/core/constants.py

@@ -18,7 +18,6 @@
 # IMPORTANT: when changing values here, update docstrings.
 
 import os
-from etils import epath
 
 # Directory in which datasets are declared within TFDS sources.
 DATASETS_TFDS_SRC_DIR = 'datasets'
@@ -48,7 +47,4 @@
 
 # Filepath for mapping between TFDS datasets and PapersWithCode entries.
 PWC_FILENAME = 'tfds_to_pwc_links.json'
-PWC_LINKS_PATH = (
-    epath.resource_path('tensorflow_datasets')
-    / f'scripts/documentation/{PWC_FILENAME}'
-)
+PWC_LINKS_PATH = f'scripts/documentation/{PWC_FILENAME}'

tensorflow_datasets/core/download/checksums.py

@@ -15,21 +15,27 @@
 
 """Methods to retrieve and store size/checksums associated to URLs."""
 
+from collections.abc import Iterable
 import dataclasses
 import hashlib
 import io
-from typing import Any, Dict, Iterable, Optional
+from typing import Any
 
 from absl import logging
 from etils import epath
 from tensorflow_datasets.core import utils
 
-_CHECKSUM_DIRS = [
-    utils.tfds_path() / 'url_checksums',
-]
+_CUSTOM_CHECKSUM_DIRS = []
 _CHECKSUM_SUFFIX = '.txt'
 
 
+@utils.memoize(maxsize=1)
+def _default_checksum_dirs() -> list[epath.Path]:
+  return [
+      utils.tfds_path() / 'url_checksums',
+  ]
+
+
 @dataclasses.dataclass(eq=True)
 class UrlInfo:
   """Small wrapper around the url metadata (checksum, size).
@@ -44,9 +50,9 @@ class UrlInfo:
   checksum: str
   # We exclude the filename from `__eq__` for backward compatibility
   # Two checksums are equals even if filename is unknown or different.
-  filename: Optional[str] = dataclasses.field(compare=False)
+  filename: str | None = dataclasses.field(compare=False)
 
-  def asdict(self) -> Dict[str, Any]:
+  def asdict(self) -> dict[str, Any]:
     """Returns the dict representation of the dataclass."""
     return dataclasses.asdict(self)
 
@@ -107,16 +113,19 @@ class MyDataset(tfds.core.DatasetBuilder):
       'The checksum file will be automatically detected. More info at: '
       'https://www.tensorflow.org/datasets/add_dataset'
   )
-  if checksums_dir in _CHECKSUM_DIRS:  # Avoid duplicate
+  if (
+      checksums_dir in _CUSTOM_CHECKSUM_DIRS
+      or checksums_dir in _default_checksum_dirs()
+  ):  # Avoid duplicates
     return
-  _CHECKSUM_DIRS.append(checksums_dir)
+  _CUSTOM_CHECKSUM_DIRS.append(checksums_dir)
 
 
 @utils.memoize()
-def _checksum_paths() -> Dict[str, epath.Path]:
+def _checksum_paths() -> dict[str, epath.Path]:
   """Returns dict {'dataset_name': 'path/to/checksums/file'}."""
   dataset2path = {}
-  for dir_path in _CHECKSUM_DIRS:
+  for dir_path in _CUSTOM_CHECKSUM_DIRS + _default_checksum_dirs():
     if isinstance(dir_path, str):
       dir_path = epath.Path(dir_path)
     if not dir_path.exists():
@@ -129,7 +138,7 @@ def _checksum_paths() -> Dict[str, epath.Path]:
   return dataset2path
 
 
-def _parse_url_infos(checksums_file: Iterable[str]) -> Dict[str, UrlInfo]:
+def _parse_url_infos(checksums_file: Iterable[str]) -> dict[str, UrlInfo]:
   """Returns {URL: (size, checksum)}s stored within given file."""
   url_infos = {}
   for line in checksums_file:
@@ -156,7 +165,7 @@ def _parse_url_infos(checksums_file: Iterable[str]) -> Dict[str, UrlInfo]:
 
 
 @utils.memoize()
-def get_all_url_infos() -> Dict[str, UrlInfo]:
+def get_all_url_infos() -> dict[str, UrlInfo]:
   """Returns dict associating URL to UrlInfo."""
   url_infos = {}
   for path in _checksum_paths().values():
@@ -171,14 +180,14 @@ def get_all_url_infos() -> Dict[str, UrlInfo]:
   return url_infos
 
 
-def load_url_infos(path: epath.PathLike) -> Dict[str, UrlInfo]:
+def load_url_infos(path: epath.PathLike) -> dict[str, UrlInfo]:
   """Loads the checksums."""
   return _parse_url_infos(epath.Path(path).read_text().splitlines())
 
 
 def save_url_infos(
     path: epath.Path,
-    url_infos: Dict[str, UrlInfo],
+    url_infos: dict[str, UrlInfo],
 ) -> None:
   """Store given checksums and sizes for specific dataset.
 
@@ -211,8 +220,8 @@ def save_url_infos(
 
 
 def _filenames_equal(
-    left: Dict[str, UrlInfo],
-    right: Dict[str, UrlInfo],
+    left: dict[str, UrlInfo],
+    right: dict[str, UrlInfo],
 ) -> bool:
   """Compare filenames."""
   return all(