Fix type errors in croissant_utils and croissant_builder.

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/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,7 +184,7 @@ 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 = base64.urlsafe_b64encode(_decode_hex(checksum))  # pyrefly: ignore[bad-assignment]
   checksum = checksum.decode()[:-1]
   name, extension = _sanitize_url(url, max_length=46)
   return f'{name}{checksum}{extension}'

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},
         ),

tensorflow_datasets/datasets/imagenet_r/imagenet_r_dataset_builder.py

@@ -60,7 +60,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_r_root': imagenet_r_root},
         ),

tensorflow_datasets/datasets/imagenet_resized/imagenet_resized_dataset_builder.py

@@ -64,7 +64,7 @@ class Builder(tfds.core.GeneratorBasedBuilder):
 
   def _info(self):
     names_file = tfds.core.tfds_path(_LABELS_FNAME)
-    size = self.builder_config.size
+    size = self.builder_config.size  # pyrefly: ignore[missing-attribute]
     return self.dataset_info_from_configs(
         features=tfds.features.FeaturesDict({
             'image': tfds.features.Image(shape=(size, size, 3)),
@@ -75,7 +75,7 @@ def _info(self):
     )
 
   def _split_generators(self, dl_manager):
-    size = self.builder_config.size
+    size = self.builder_config.size  # pyrefly: ignore[missing-attribute]
 
     if size in [8, 16, 32]:
       train_path, val_path = dl_manager.download([
@@ -96,7 +96,7 @@ def _split_generators(self, dl_manager):
 
     return [
         tfds.core.SplitGenerator(
-            name=tfds.Split.TRAIN,
+            name=tfds.Split.TRAIN,  # pyrefly: ignore[missing-attribute]
             gen_kwargs={
                 'archive': itertools.chain(
                     *[
@@ -107,7 +107,7 @@ def _split_generators(self, dl_manager):
             },
         ),
         tfds.core.SplitGenerator(
-            name=tfds.Split.VALIDATION,
+            name=tfds.Split.VALIDATION,  # pyrefly: ignore[missing-attribute]
             gen_kwargs={
                 'archive': dl_manager.iter_archive(val_path),
             },
@@ -121,7 +121,7 @@ def _generate_examples(self, archive):
       if content:
         fobj_mem = io.BytesIO(content)
         data = np.load(fobj_mem, allow_pickle=False)
-        size = self.builder_config.size
+        size = self.builder_config.size  # pyrefly: ignore[missing-attribute]
         for i, (image, label) in enumerate(zip(data['data'], data['labels'])):
           record = {
               # The data is packed flat as CHW where as most image datasets

tensorflow_datasets/datasets/imagenet_sketch/imagenet_sketch_dataset_builder.py

@@ -48,7 +48,7 @@ def _split_generators(self, dl_manager: tfds.download.DownloadManager):
     path = dl_manager.download(_IMAGENET_SKETCH_URL)
     return [
         tfds.core.SplitGenerator(
-            name=tfds.Split.TEST,
+            name=tfds.Split.TEST,  # pyrefly: ignore[missing-attribute]
             gen_kwargs={
                 'archive': dl_manager.iter_archive(path),
             },