Fix lsun tests on GitHub by importing lmdb directly.

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

tensorflow_datasets/core/utils/croissant_utils.py

@@ -213,7 +213,10 @@ def get_split_recordset(
       if not record_sets:
         raise ValueError(f"Field {field.id} has no RecordSet.")
       referenced_record_set = record_sets[0]
-      if mlc.DataType.SPLIT in referenced_record_set.data_types:
+      if (
+          referenced_record_set.data_types
+          and mlc.DataType.SPLIT in referenced_record_set.data_types
+      ):
         return SplitReference(referenced_record_set, field)
   return None
 
@@ -233,3 +236,5 @@ def get_record_set_ids(metadata: mlc.Metadata) -> list[str]:
       continue
     record_set_ids.append(record_set.id)
   return record_set_ids
+
+# Dummy comment to force pytype run

tensorflow_datasets/datasets/dsprites/dsprites_dataset_builder.py

@@ -62,7 +62,7 @@ def _split_generators(self, dl_manager):
     # There is no predefined train/val/test split for this dataset.
     return [
         tfds.core.SplitGenerator(
-            name=tfds.Split.TRAIN, gen_kwargs=dict(filepath=filepath)
+            name=tfds.Split.TRAIN, gen_kwargs=dict(filepath=filepath)  # pyrefly: ignore[missing-attribute]
         ),
     ]
 

tensorflow_datasets/datasets/e2e_cleaned/e2e_cleaned_dataset_builder.py

@@ -77,17 +77,17 @@ def _split_generators(self, dl_manager):
     )
     return [
         tfds.core.SplitGenerator(
-            name=tfds.Split.TRAIN,
+            name=tfds.Split.TRAIN,  # pyrefly: ignore[missing-attribute]
             # These kwargs will be passed to _generate_examples
             gen_kwargs={'csv_path': extracted_path['train_path']},
         ),
         tfds.core.SplitGenerator(
-            name=tfds.Split.VALIDATION,
+            name=tfds.Split.VALIDATION,  # pyrefly: ignore[missing-attribute]
             # These kwargs will be passed to _generate_examples
             gen_kwargs={'csv_path': extracted_path['dev_path']},
         ),
         tfds.core.SplitGenerator(
-            name=tfds.Split.TEST,
+            name=tfds.Split.TEST,  # pyrefly: ignore[missing-attribute]
             # These kwargs will be passed to _generate_examples
             gen_kwargs={'csv_path': extracted_path['test_path']},
         ),

tensorflow_datasets/datasets/efron_morris75/efron_morris75_dataset_builder.py

@@ -54,7 +54,7 @@ def _info(self) -> tfds.core.DatasetInfo:
   def _split_generators(self, dl_manager: tfds.download.DownloadManager):
     """Returns SplitGenerators."""
     path = dl_manager.download_and_extract(URL)
-    return {tfds.Split.TRAIN: self._generate_examples(path)}
+    return {tfds.Split.TRAIN: self._generate_examples(path)}  # pyrefly: ignore[missing-attribute]
 
   def _generate_examples(self, path):
     """Yields examples."""

tensorflow_datasets/datasets/flic/flic_dataset_builder.py

@@ -99,10 +99,10 @@ def _info(self):
 
   def _split_generators(self, dl_manager):
     """Returns SplitGenerators."""
-    extract_path = dl_manager.download_and_extract(self.builder_config.url)
+    extract_path = dl_manager.download_and_extract(self.builder_config.url)  # pyrefly: ignore[missing-attribute]
 
     mat_path = os.path.join(
-        extract_path, self.builder_config.dir, "examples.mat"
+        extract_path, self.builder_config.dir, "examples.mat"  # pyrefly: ignore[missing-attribute]
     )
     with tf.io.gfile.GFile(mat_path, "rb") as f:
       data = tfds.core.lazy_imports.scipy.io.loadmat(
@@ -111,15 +111,15 @@ def _split_generators(self, dl_manager):
 
     return [
         tfds.core.SplitGenerator(
-            name=tfds.Split.TRAIN,
+            name=tfds.Split.TRAIN,  # pyrefly: ignore[missing-attribute]
             gen_kwargs={
                 "extract_path": extract_path,
                 "data": data,
                 "selection_column": 7,  # indicates train split selection
             },
         ),
         tfds.core.SplitGenerator(
-            name=tfds.Split.TEST,
+            name=tfds.Split.TEST,  # pyrefly: ignore[missing-attribute]
             gen_kwargs={
                 "extract_path": extract_path,
                 "data": data,
@@ -133,7 +133,7 @@ def _generate_examples(self, extract_path, data, selection_column):
     for u_id, example in enumerate(data["examples"]):
       if example[selection_column]:
         img_path = os.path.join(
-            extract_path, self.builder_config.dir, "images", example[3]
+            extract_path, self.builder_config.dir, "images", example[3]  # pyrefly: ignore[missing-attribute]
         )
         yield u_id, {
             "image": img_path,

tensorflow_datasets/datasets/groove/groove_dataset_builder.py

@@ -130,12 +130,12 @@ def _info(self):
         },
         "midi": tf.string,
     }
-    if self.builder_config.include_audio:
-      features_dict["audio"] = tfds.features.Audio(
-          dtype=np.float32, sample_rate=self.builder_config.audio_rate
+    if self.builder_config.include_audio:  # pyrefly: ignore[missing-attribute]
+      features_dict["audio"] = tfds.features.Audio(  # pyrefly: ignore[bad-assignment]
+          dtype=np.float32, sample_rate=self.builder_config.audio_rate  # pyrefly: ignore[missing-attribute]
       )
     return self.dataset_info_from_configs(
-        features=tfds.features.FeaturesDict(features_dict),
+        features=tfds.features.FeaturesDict(features_dict),  # pyrefly: ignore[bad-argument-type]
         homepage="https://g.co/magenta/groove-dataset",
     )
 

tensorflow_datasets/datasets/imagenet2012/imagenet2012_dataset_builder.py

@@ -101,20 +101,20 @@ def _split_generators(self, dl_manager):
     splits = []
     _add_split_if_exists(
         split_list=splits,
-        split=tfds.Split.TRAIN,
+        split=tfds.Split.TRAIN,  # pyrefly: ignore[missing-attribute]
         split_path=train_path,
         dl_manager=dl_manager,
     )
     _add_split_if_exists(
         split_list=splits,
-        split=tfds.Split.VALIDATION,
+        split=tfds.Split.VALIDATION,  # pyrefly: ignore[missing-attribute]
         split_path=val_path,
         dl_manager=dl_manager,
         validation_labels=imagenet_common.get_validation_labels(val_path),
     )
     _add_split_if_exists(
         split_list=splits,
-        split=tfds.Split.TEST,
+        split=tfds.Split.TEST,  # pyrefly: ignore[missing-attribute]
         split_path=test_path,
         dl_manager=dl_manager,
         labels_exist=False,

tensorflow_datasets/datasets/imagenet2012_corrupted/imagenet2012_corrupted_dataset_builder.py

@@ -168,7 +168,7 @@ def _split_generators(self, dl_manager):
     splits = super(Builder, self)._split_generators(dl_manager)
 
     corruptions.FROST_FILENAMES = dl_manager.download(_FROST_FILENAMES)
-    return [s for s in splits if s.name != tfds.Split.TRAIN]
+    return [s for s in splits if s.name != tfds.Split.TRAIN]  # pyrefly: ignore[missing-attribute]
 
   def _generate_examples(
       self, archive, validation_labels=None, labels_exist=None
@@ -215,8 +215,8 @@ def _get_corrupted_example(self, x):
     Returns:
       numpy array, corrupted images.
     """
-    corruption_type = self.builder_config.corruption_type
-    severity = self.builder_config.severity
+    corruption_type = self.builder_config.corruption_type  # pyrefly: ignore[missing-attribute]
+    severity = self.builder_config.severity  # pyrefly: ignore[missing-attribute]
     x = np.clip(x, 0, 255)
 
     return {

tensorflow_datasets/datasets/imagenet2012_fewshot/imagenet2012_fewshot_dataset_builder.py

@@ -68,7 +68,7 @@ def _split_generators(self, dl_manager):
       )
 
     # Download and load subset file.
-    subset_file = SUBSET2FILES[self.builder_config.name]
+    subset_file = SUBSET2FILES[self.builder_config.name]  # pyrefly: ignore[missing-attribute]
     if isinstance(subset_file, list):  # it will only be a list during testing,
       subset_file = subset_file[0]  # where the first entry is 1shot.txt.
     subset = set(subset_file.read_text().splitlines())
@@ -77,13 +77,13 @@ def _split_generators(self, dl_manager):
     tuneset = set(TUNE_FILE.read_text().splitlines())
 
     return {
-        tfds.Split.TRAIN: self._generate_examples(
+        tfds.Split.TRAIN: self._generate_examples(  # pyrefly: ignore[missing-attribute]
             archive=dl_manager.iter_archive(train_path), subset=subset
         ),
         tfds.Split('tune'): self._generate_examples(
             archive=dl_manager.iter_archive(train_path), subset=tuneset
         ),
-        tfds.Split.VALIDATION: self._generate_examples(
+        tfds.Split.VALIDATION: self._generate_examples(  # pyrefly: ignore[missing-attribute]
             archive=dl_manager.iter_archive(val_path),
             validation_labels=imagenet_common.get_validation_labels(val_path),
         ),

tensorflow_datasets/datasets/imagenet2012_real/imagenet2012_real_dataset_builder.py

@@ -72,7 +72,7 @@ def _split_generators(self, dl_manager):
       )
     return [
         tfds.core.SplitGenerator(
-            name=tfds.Split.VALIDATION,
+            name=tfds.Split.VALIDATION,  # pyrefly: ignore[missing-attribute]
             gen_kwargs={
                 'archive': dl_manager.iter_archive(val_path),
                 'original_labels': imagenet_common.get_validation_labels(

tensorflow_datasets/datasets/imagenet2012_subset/imagenet2012_subset_dataset_builder.py

@@ -72,22 +72,22 @@ def _split_generators(self, dl_manager):
       )
 
     # Download and load subset file.
-    subset_file = dl_manager.download(SUBSET2FILES[self.builder_config.name])
+    subset_file = dl_manager.download(SUBSET2FILES[self.builder_config.name])  # pyrefly: ignore[missing-attribute]
     if isinstance(subset_file, list):  # it will only be a list during testing,
       subset_file = subset_file[0]  # where the first entry is 1percent.txt.
     with epath.Path(subset_file).open() as fp:
       subset = set(fp.read().splitlines())  # remove trailing `\r` in Windows
 
     return [
         tfds.core.SplitGenerator(
-            name=tfds.Split.TRAIN,
+            name=tfds.Split.TRAIN,  # pyrefly: ignore[missing-attribute]
             gen_kwargs={
                 'archive': dl_manager.iter_archive(train_path),
                 'subset': subset,
             },
         ),
         tfds.core.SplitGenerator(
-            name=tfds.Split.VALIDATION,
+            name=tfds.Split.VALIDATION,  # pyrefly: ignore[missing-attribute]
             gen_kwargs={
                 'archive': dl_manager.iter_archive(val_path),
                 'validation_labels': imagenet_common.get_validation_labels(

tensorflow_datasets/datasets/imagenet_a/imagenet_a_dataset_builder.py

@@ -51,7 +51,7 @@ def _split_generators(self, dl_manager):
     return [
         tfds.core.SplitGenerator(
             # The dataset provides only a test split.
-            name=tfds.Split.TEST,
+            name=tfds.Split.TEST,  # pyrefly: ignore[missing-attribute]
             # These kwargs will be passed to _generate_examples
             gen_kwargs={'imagenet_a_root': imagenet_a_root},
         ),