Fix lsun tests on GitHub by importing lmdb directly. (#11232)

PiperOrigin-RevId: 968380231

Co-authored-by: Pierre Marcenac <pierremarcenac@google.com>

Author: copybara-service[bot]

setup.py

@@ -153,6 +153,7 @@
         # tensorflow==2.20.0 is not compatible with gcld3 because of protobuf
         # version conflict.
         'tensorflow<2.20.0',
+        'tensorflow-metadata<1.18.0',
     ],
     'c4_wsrs': ['apache-beam<2.65.0'],
     'cats_vs_dogs': ['matplotlib'],
@@ -170,16 +171,10 @@
         'scipy',
     ],
     'librispeech': ['pydub'],  # and ffmpeg installed
-    'lsun': [
-        # tensorflow-io is compiled against specific versions of TF.
-        'tensorflow-io[tensorflow]',
-    ],
+    'lsun': ['lmdb'],
     'nsynth': [
         'crepe',
         'librosa',
-        # tensorflow==2.20.0 is not compatible with librosa because of protobuf
-        # version conflict.
-        'tensorflow<2.20.0',
     ],
     'ogbg_molpcba': ['pandas', 'networkx'],
     'pet_finder': ['pandas'],

tensorflow_datasets/core/community/register_package_test.py

@@ -105,7 +105,7 @@ def test_builder_cls(dummy_register: register_package.PackageRegister):  # pylin
   builder_cls = dummy_register.builder_cls(ds_name)
   assert builder_cls.name == 'dummy_dataset'
 
-  clshash = '89a6603e427554b46d2b4895b24ce8d4c2d703d453deaff3946c941fefaf1412'
+  clshash = '925292f8861ff81947528d2ed011760c47e7a09a9b29062148e81d7507fc4e60'
   assert installed_path / f'{clshash}/dummy_dataset.py' == builder_cls.code_path
   assert 'kaggle' in builder_cls.code_path.parts
   assert issubclass(builder_cls, dataset_builder.DatasetBuilder)

tensorflow_datasets/core/dataset_builders/croissant_builder.py

@@ -39,7 +39,7 @@
 from collections.abc import Mapping, Sequence
 import datetime
 import json
-from typing import Any
+from typing import Any, cast
 
 from etils import enp
 from etils import epath
@@ -110,7 +110,7 @@ def array_datatype_converter(
   if field.data_type in dtype_mapping:
     field_dtype = dtype_mapping[field.data_type]
   elif enp.lazy.is_np_dtype(field.data_type):
-    field_dtype = field.data_type
+    field_dtype = cast(type_utils.TfdsDType, field.data_type)
 
   description = croissant_utils.extract_localized_string(
       field.description, language=language, field_name='description'
@@ -191,7 +191,7 @@ def datatype_converter(
   elif field_data_type in dtype_mapping:
     feature = dtype_mapping[field_data_type]
   elif enp.lazy.is_np_dtype(field_data_type):
-    feature = field_data_type
+    feature = cast(type_utils.TfdsDType, field_data_type)
   # We return a text feature for date-time features (mlc.DataType.DATE,
   # mlc.DataType.DATETIME, and mlc.DataType.TIME).
   elif field_data_type == pd.Timestamp or field_data_type == datetime.time:

tensorflow_datasets/core/dataset_info.py

@@ -1007,7 +1007,7 @@ def _nest_to_proto(nest: Nest) -> dataset_info_pb2.SupervisedKeys.Nest:
     for item in nest:
       proto.tuple.items.append(_nest_to_proto(item))
   elif nest_type is dict:
-    nest = {key: _nest_to_proto(value) for key, value in nest.items()}  # pyrefly: ignore[bad-assignment]
+    nest = {key: _nest_to_proto(value) for key, value in nest.items()}  # pyrefly: ignore[bad-assignment, missing-attribute]
     proto.dict.CopyFrom(dataset_info_pb2.SupervisedKeys.Dict(dict=nest))  # pyrefly: ignore[bad-argument-type]
   elif nest_type is str:
     proto.feature_key = nest  # pyrefly: ignore[bad-assignment]

tensorflow_datasets/core/download/download_manager.py

@@ -348,7 +348,7 @@ def register_checksums(self):
   def _record_url_infos(self):
     """Store in file when recorded size/checksum of downloaded files."""
     checksums.save_url_infos(
-        self._register_checksums_path,
+        self._register_checksums_path,  # pyrefly: ignore[bad-argument-type]
         self._recorded_url_infos,
     )
 
@@ -575,7 +575,7 @@ def callback(dl_result: downloader.DownloadResult) -> epath.Path:
       resource_lib.write_info_file(
           url=url,
           path=dst_path,
-          dataset_name=self._dataset_name,
+          dataset_name=self._dataset_name,  # pyrefly: ignore[bad-argument-type]
           original_fname=dl_path.name,
           url_info=dl_url_info,
       )

tensorflow_datasets/core/download/downloader.py

@@ -128,7 +128,7 @@ def _get_filename(response: Response) -> str:
     if filename:
       return filename
   # Otherwise, fallback on extracting the name from the url.
-  return _basename_from_url(response.url)
+  return _basename_from_url(response.url)  # pyrefly: ignore[bad-argument-type]
 
 
 def _process_gdrive_confirmation(original_url: str, contents: str) -> str:
@@ -345,7 +345,7 @@ def _open_with_requests(
 ) -> Iterator[tuple[Response, Iterable[bytes]]]:
   """Open url with request."""
   with requests.Session() as session:
-    retries = requests.packages.urllib3.util.retry.Retry(
+    retries = requests.packages.urllib3.util.retry.Retry(  # pyrefly: ignore[missing-attribute]
         total=MAX_RETRIES,
         backoff_factor=0.2,
         status_forcelist=[500, 502, 503, 504],

tensorflow_datasets/core/download/extractor.py

@@ -184,7 +184,7 @@ def iter_tar(arch_f, stream=False):
   read_type = 'r' + ('|' if stream else ':') + '*'
 
   with _open_or_pass(arch_f) as fobj:
-    tar = tarfile.open(mode=read_type, fileobj=fobj)
+    tar = tarfile.open(mode=read_type, fileobj=fobj)  # pyrefly: ignore[no-matching-overload]
     for member in tar:
       if stream and (member.islnk() or member.issym()):
         # Links cannot be dereferenced in stream mode.

tensorflow_datasets/core/download/resource.py

@@ -130,15 +130,15 @@ def _sanitize_url(url: str, max_length: int) -> tuple[str, str]:
   Returns:
     Sanitized and shorted url, file extension.
   """
-  url = urllib.parse.urlparse(url)
-  netloc = url.netloc
+  url = urllib.parse.urlparse(url)  # pyrefly: ignore[bad-assignment]
+  netloc = url.netloc  # pyrefly: ignore[missing-attribute]
   for prefix in _NETLOC_COMMON_PREFIXES:
     if netloc.startswith(prefix):
       netloc = netloc[len(prefix) :]
   for suffix in _NETLOC_COMMON_SUFFIXES:
     if netloc.endswith(suffix):
       netloc = netloc[: -len(suffix)]
-  url = f'{netloc}{url.path}{url.params}{url.query}'
+  url = f'{netloc}{url.path}{url.params}{url.query}'  # pyrefly: ignore[missing-attribute]
   # Get the extension:
   for ext in _KNOWN_EXTENSIONS:
     if url.endswith(ext):
@@ -184,8 +184,8 @@ def get_dl_fname(url: str, checksum: str | None = None) -> str:
   """
   if not checksum:
     checksum = checksums_lib.sha256(url)
-  checksum = base64.urlsafe_b64encode(_decode_hex(checksum))
-  checksum = checksum.decode()[:-1]
+  checksum = base64.urlsafe_b64encode(_decode_hex(checksum))  # pyrefly: ignore[bad-assignment]
+  checksum = checksum.decode()[:-1]  # pyrefly: ignore[missing-attribute]
   name, extension = _sanitize_url(url, max_length=46)
   return f'{name}{checksum}{extension}'
 

tensorflow_datasets/core/lazy_imports_lib.py

@@ -103,6 +103,11 @@ def langdetect(cls):
   def librosa(cls):
     return _try_import("librosa")
 
+  @utils.classproperty
+  @classmethod
+  def lmdb(cls):
+    return _try_import("lmdb")
+
   @utils.classproperty
   @classmethod
   def lxml(cls):

tensorflow_datasets/core/read_only_builder.py

@@ -245,7 +245,7 @@ def get_split_dict(builder_dir, dataset_info_proto) -> splits_lib.SplitDict:
 
   merged_split_dict = splits_lib.SplitDict.merge_multiple([
       get_split_dict(builder_dir, dataset_info_proto)
-      for builder_dir, dataset_info_proto in dataset_infos.items()
+      for builder_dir, dataset_info_proto in dataset_infos.items()  # pyrefly: ignore[missing-attribute]
   ])
 
   # We create the ReadOnlyBuilder for a random builder_dir and then update the