text: # coding=utf-8
  # Copyright 2021 The TensorFlow Datasets Authors.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  
  """Base DatasetBuilderTestCase to test a DatasetBuilder base class."""
  
  import collections
  import contextlib
  import difflib
  import hashlib
  import itertools
  import numbers
  import os
  import textwrap
  from typing import Iterator
  from unittest import mock
  
  from absl.testing import parameterized
  import numpy as np
  import tensorflow.compat.v2 as tf
  
  from tensorflow_datasets.core import dataset_builder
  from tensorflow_datasets.core import dataset_info
  from tensorflow_datasets.core import dataset_utils
  from tensorflow_datasets.core import download
  from tensorflow_datasets.core import load
  from tensorflow_datasets.core import utils
  from tensorflow_datasets.core.download import checksums
  from tensorflow_datasets.testing import feature_test_case
  from tensorflow_datasets.testing import test_utils
  
  # `os` module Functions for which tf.io.gfile equivalent should be preferred.
  FORBIDDEN_OS_FUNCTIONS = (
      "chmod",
      "chown",
      "link",
      "listdir",
      "lstat",
      "makedirs",
      "mkdir",
      "mknod",
      "open",
      "pathconf",
      "readlink",
      "remove",
      "removedirs",
      "rename",
      "renames",
      "rmdir",
      "stat",
      "statvfs",
      "symlink",
      "unlink",
      "walk",
  )
  FORBIDDEN_OS_PATH_FUNCTIONS = (
      "exists",
      "isdir",
      "isfile",
  )
  
  
  _ORGINAL_NP_LOAD = np.load
  
  
  def _np_load(file_, mmap_mode=None, allow_pickle=False, **kwargs):
    if not hasattr(file_, "read"):
      raise AssertionError(
          "You MUST pass a `tf.io.gfile.GFile` or file-like object to `np.load`.")
    if allow_pickle:
      raise AssertionError("Unpicling files is forbidden for security reasons.")
    return _ORGINAL_NP_LOAD(file_, mmap_mode, allow_pickle, **kwargs)
  
  
  class DatasetBuilderTestCase(
      parameterized.TestCase, feature_test_case.SubTestCase
  ):
    """Inherit this class to test your DatasetBuilder class.
  
    You must set the following class attributes:
  
      * DATASET_CLASS: class object of DatasetBuilder you want to test.
  
    You may set the following class attributes:
  
      * VERSION: `str`. The version used to run the test. eg: '1.2.*'.
        Defaults to None (canonical version).
      * BUILDER_CONFIG_NAMES_TO_TEST: `list[str]`, the list of builder configs
        that should be tested. If None, all the BUILDER_CONFIGS from the class
        will be tested.
      * DL_EXTRACT_RESULT: `dict[str]`, the returned result of mocked
        `download_and_extract` method. The values should be the path of files
        present in the `fake_examples` directory, relative to that directory.
        If not specified, path to `fake_examples` will always be returned.
      * DL_DOWNLOAD_RESULT: `dict[str]`, the returned result of mocked
        `download_and_extract` method. The values should be the path of files
        present in the `fake_examples` directory, relative to that directory.
        If not specified: will use DL_EXTRACT_RESULT (this is due to backwards
        compatibility and will be removed in the future).
      * EXAMPLE_DIR: `str`, the base directory in in which fake examples are
        contained. Optional; defaults to
        tensorflow_datasets/testing/test_data/fake_examples/<dataset name>.
      * OVERLAPPING_SPLITS: `list[str]`, splits containing examples from other
        splits (e.g. a "example" split containing pictures from other splits).
      * MOCK_OUT_FORBIDDEN_OS_FUNCTIONS: `bool`, defaults to True. Set to False to
        disable checks preventing usage of `os` or builtin functions instead of
        recommended `tf.io.gfile` API.
      * SKIP_CHECKSUMS: Checks that the urls called by `dl_manager.download`
        are registered.
  
    This test case will check for the following:
  
     - the dataset builder is correctly registered, i.e. `tfds.load(name)` works;
     - the dataset builder can read the fake examples stored in
         testing/test_data/fake_examples/{dataset_name};
     - the dataset builder can produce serialized data;
     - the dataset builder produces a valid Dataset object from serialized data
       - in eager mode;
       - in graph mode.
     - the produced Dataset examples have the expected dimensions and types;
     - the produced Dataset has and the expected number of examples;
     - a example is not part of two splits, or one of these splits is whitelisted
         in OVERLAPPING_SPLITS.
    """
  
    DATASET_CLASS = None
    VERSION = None
    BUILDER_CONFIG_NAMES_TO_TEST = None
    DL_EXTRACT_RESULT = None
    DL_DOWNLOAD_RESULT = None
    EXAMPLE_DIR = None
    OVERLAPPING_SPLITS = []
    MOCK_OUT_FORBIDDEN_OS_FUNCTIONS = True
    SKIP_CHECKSUMS = False
  
    @classmethod
    def setUpClass(cls):
      tf.enable_v2_behavior()
      super(DatasetBuilderTestCase, cls).setUpClass()
      name = cls.__name__
      # Check class has the right attributes
      if cls.DATASET_CLASS is None or not callable(cls.DATASET_CLASS):
        raise AssertionError(
            "Assign your DatasetBuilder class to %s.DATASET_CLASS." % name)
  
    def setUp(self):
      super(DatasetBuilderTestCase, self).setUp()
      self.patchers = []
      self.builder = self._make_builder()
  
      example_dir = self.DATASET_CLASS.code_path.parent / "dummy_data"
      fake_example_dir = utils.as_path(test_utils.fake_examples_dir())
      if self.EXAMPLE_DIR is not None:
        self.example_dir = utils.as_path(self.EXAMPLE_DIR)
        example_dir = self.example_dir  # Dir to display in the error
      elif example_dir.exists():
        self.example_dir = example_dir
      else:
        self.example_dir = fake_example_dir / self.builder.name
  
      if not self.example_dir.exists():
        err_msg = f"Dummy data not found in: {example_dir}"
        raise ValueError(err_msg)
      if self.MOCK_OUT_FORBIDDEN_OS_FUNCTIONS:
        self._mock_out_forbidden_os_functions()
  
      # Track the urls which are downloaded to validate the checksums
      # The `dl_manager.download` and `dl_manager.download_and_extract` are
      # patched to record the urls in `_download_urls`.
      # Calling `dl_manager.download_checksums` stop the url
      # registration (as checksums are stored remotelly)
      # `_test_checksums` validates the recorded urls.
      self._download_urls = set()
      self._stop_record_download = False
  
    def tearDown(self):
      super(DatasetBuilderTestCase, self).tearDown()
      for patcher in self.patchers:
        patcher.stop()
  
    def _mock_out_forbidden_os_functions(self):
      """Raises error if forbidden os functions are called instead of gfile."""
      err = AssertionError("Do not use `os`, but `tf.io.gfile` module instead. "
                           "This makes code compatible with more filesystems.")
      sep = os.path.sep
      mock_os_path = mock.Mock(os.path, wraps=os.path)
      mock_os_path.sep = sep
      for fop in FORBIDDEN_OS_PATH_FUNCTIONS:
        getattr(mock_os_path, fop).side_effect = err
      mock_os = mock.Mock(os, path=mock_os_path, fspath=os.fspath)
      for fop in FORBIDDEN_OS_FUNCTIONS:
        if os.name == "nt" and not hasattr(os, fop):
          continue  # Not all `os` functions are available on Windows (ex: chmod).
        getattr(mock_os, fop).side_effect = err
      os_patcher = mock.patch(
          self.DATASET_CLASS.__module__ + ".os", mock_os, create=True)
      os_patcher.start()
      self.patchers.append(os_patcher)
  
      mock_builtins = __builtins__.copy()  # pytype: disable=module-attr
      mock_builtins["open"] = mock.Mock(side_effect=err)
      open_patcher = mock.patch(
          self.DATASET_CLASS.__module__ + ".__builtins__", mock_builtins)
      open_patcher.start()
      self.patchers.append(open_patcher)
  
      # It's hard to mock open within numpy, so mock np.load.
      np_load_patcher = mock.patch("numpy.load", _np_load)
      np_load_patcher.start()
      self.patchers.append(np_load_patcher)
  
    def test_baseclass(self):
      self.assertIsInstance(
          self.builder, dataset_builder.DatasetBuilder,
          "Dataset class must inherit from `dataset_builder.DatasetBuilder`.")
      # Since class was instantiated and base class is ABCMeta, then we know
      # all needed methods were implemented.
  
    def test_registered(self):
      is_registered = self.builder.name in load.list_builders(
          with_community_datasets=False,
      )
      self.assertTrue(
          is_registered,
          f"Dataset {self.builder.name} was not registered.",
      )
  
    def test_info(self):
      info = self.builder.info
      self.assertIsInstance(info, dataset_info.DatasetInfo)
      self.assertEqual(self.builder.name, info.name)
  
    def _add_url(self, url_or_urls):
      if self._stop_record_download:
        # Stop record the checksums if dl_manager.download_checksums has been
        # called (as checksums may be stored remotelly)
        return
      if isinstance(url_or_urls, download.resource.Resource):
        self._download_urls.add(url_or_urls.url)
      else:
        self._download_urls.add(url_or_urls)
  
    def _get_dl_extract_result(self, url):
      tf.nest.map_structure(self._add_url, url)
      del url
      if self.DL_EXTRACT_RESULT is None:
        return self.example_dir
      return tf.nest.map_structure(
          lambda fname: self.example_dir / fname,
          self.DL_EXTRACT_RESULT,
      )
  
    def _get_dl_download_result(self, url):
      tf.nest.map_structure(self._add_url, url)
      if self.DL_DOWNLOAD_RESULT is None:
        # This is only to be backwards compatible with old approach.
        # In the future it will be replaced with using self.example_dir.
        return self._get_dl_extract_result(url)
      return tf.nest.map_structure(
          lambda fname: self.example_dir / fname,
          self.DL_DOWNLOAD_RESULT,
      )
  
    def _download_checksums(self, url):
      self._stop_record_download = True
  
    def _make_builder(self, config=None):
      return self.DATASET_CLASS(  # pylint: disable=not-callable
          data_dir=self.tmp_dir,
          config=config,
          version=self.VERSION)
  
    @test_utils.run_in_graph_and_eager_modes()
    def test_download_and_prepare_as_dataset(self):
      # If configs specified, ensure they are all valid
      if self.BUILDER_CONFIG_NAMES_TO_TEST:
        for config in self.BUILDER_CONFIG_NAMES_TO_TEST:  # pylint: disable=not-an-iterable
          assert config in self.builder.builder_configs, (
              "Config %s specified in test does not exist. Available:\n%s" % (
                  config, list(self.builder.builder_configs)))
  
      configs = self.builder.BUILDER_CONFIGS
      print("Total configs: %d" % len(configs))
      if configs:
        for config in configs:
          # Skip the configs that are not in the list.
          if (self.BUILDER_CONFIG_NAMES_TO_TEST is not None and
              (config.name not in self.BUILDER_CONFIG_NAMES_TO_TEST)):  # pylint: disable=unsupported-membership-test
            print("Skipping config %s" % config.name)
            continue
          with self._subTest(config.name):
            print("Testing config %s" % config.name)
            builder = self._make_builder(config=config)
            self._download_and_prepare_as_dataset(builder)
      else:
        self._download_and_prepare_as_dataset(self.builder)
  
      if not self.SKIP_CHECKSUMS:
        with self._subTest("url_checksums"):
          self._test_checksums()
  
    def _test_checksums(self):
      # If no call to `dl_manager.download`, then no need to check url presence.
      if not self._download_urls:
        return
  
      err_msg = (
          "Did you forget to record checksums with `--register_checksums` ? "
          "See instructions at: "
          "hhttps://www.tensorflow.org/datasets/add_dataset#run_the_generation_code"
          "If want to opt-out of checksums validation, please add "
          "`SKIP_CHECKSUMS = True` to the `DatasetBuilderTestCase`.\n"
      )
      url_infos = self.DATASET_CLASS.url_infos
      filepath = self.DATASET_CLASS._checksums_path  # pylint: disable=protected-access
      # Legacy checksums: Search in `checksums/` dir
      if url_infos is None:
        legacy_filepath = checksums._checksum_paths().get(self.builder.name)  # pylint: disable=protected-access
        if legacy_filepath and legacy_filepath.exists():
          filepath = legacy_filepath
          url_infos = checksums.load_url_infos(filepath)
      # Checksums not present neither in legacy nor package
      if url_infos is None:
        raise FileNotFoundError(
            f"Checksums file not found at: {filepath}\n"
            f"{err_msg}\n"
        )
  
      missing_urls = self._download_urls - set(url_infos.keys())
      self.assertEmpty(
          missing_urls,
          f"Some urls checksums are missing at: {filepath}\n{err_msg}"
      )
  
    def _download_and_prepare_as_dataset(self, builder):
      # Provide the manual dir only if builder has MANUAL_DOWNLOAD_INSTRUCTIONS
      # set.
  
      missing_dir_mock = mock.PropertyMock(
          side_effect=Exception("Missing MANUAL_DOWNLOAD_INSTRUCTIONS"))
  
      manual_dir = (
          self.example_dir
          if builder.MANUAL_DOWNLOAD_INSTRUCTIONS else missing_dir_mock)
      with mock.patch.multiple(
          "tensorflow_datasets.core.download.DownloadManager",
          download_and_extract=self._get_dl_extract_result,
          download=self._get_dl_download_result,
          download_checksums=self._download_checksums,
          manual_dir=manual_dir,
          download_dir=self.example_dir,
      ):
        # For Beam datasets, set-up the runner config
        beam_runner = None
  
        download_config = download.DownloadConfig(
            compute_stats=download.ComputeStatsMode.SKIP,
            beam_runner=beam_runner,
        )
        with self._test_key_not_local_path(builder):
          builder.download_and_prepare(download_config=download_config)
  
      with self._subTest("as_dataset"):
        self._assertAsDataset(builder)
  
      with self._subTest("num_examples"):
        self._assertNumSamples(builder)
  
      with self._subTest("reload"):
        # When reloading the dataset, metadata should been reloaded too.
  
        builder_reloaded = self._make_builder(config=builder.builder_config)
        self._assertNumSamples(builder_reloaded)
  
        # After reloading, as_dataset should still be working
        with self._subTest("as_dataset"):
          self._assertAsDataset(builder_reloaded)
  
      with self._subTest("config_description"):
        self._test_description_builder_config(builder)
  
    @contextlib.contextmanager
    def _test_key_not_local_path(self, builder) -> Iterator[None]:
      if not isinstance(builder, dataset_builder.GeneratorBasedBuilder):
        yield  # Do not test non-generator builder
        return
  
      original_generate_examples = builder._generate_examples  # pylint: disable=protected-access
  
      def _iter_examples(generator):
        for key, ex in generator:
          self._assert_key_valid(key)
          yield key, ex
  
      def new_generate_examples(*args, **kwargs):
        examples = original_generate_examples(*args, **kwargs)
        try:
          import apache_beam as beam  # pylint: disable=g-import-not-at-top
        except ImportError:
          beam = None
        if beam and isinstance(examples, (beam.PCollection, beam.PTransform)):
          return examples  # Beam keys not supported for now
        elif isinstance(examples, collections.abc.Iterable):
          return _iter_examples(examples)
        else:  # Unexpected
          return examples
  
      with mock.patch.object(
          builder, "_generate_examples", new_generate_examples
      ):
        yield
  
    def _assert_key_valid(self, key):
      if isinstance(key, str) and os.fspath(self.example_dir) in key:
        err_msg = (
            "Key yield in '_generate_examples' method "
            f"contain user directory path: {key}.\n"
            "This makes the dataset example order non-deterministic. "
            "Please use `filepath.name`, or `os.path.basename` instead."
        )
        raise ValueError(err_msg)
  
    def _assertAsDataset(self, builder):
      split_to_checksums = {}  # {"split": set(examples_checksums)}
      for split_name, expected_examples_number in self.SPLITS.items():
        ds = builder.as_dataset(split=split_name)
        compare_shapes_and_types(
            builder.info.features.get_tensor_info(),
            tf.compat.v1.data.get_output_types(ds),
            tf.compat.v1.data.get_output_shapes(ds),
        )
        examples = list(dataset_utils.as_numpy(
            builder.as_dataset(split=split_name)))
        split_to_checksums[split_name] = set(checksum(rec) for rec in examples)
        self.assertLen(examples, expected_examples_number)
      for (split1, hashes1), (split2, hashes2) in itertools.combinations(
          split_to_checksums.items(), 2):
        if (split1 in self.OVERLAPPING_SPLITS or
            split2 in self.OVERLAPPING_SPLITS):
          continue
        self.assertFalse(
            hashes1.intersection(hashes2),
            ("Splits '%s' and '%s' are overlapping. Are you sure you want to "
             "have the same objects in those splits? If yes, add one one of "
             "them to OVERLAPPING_SPLITS class attribute.") % (split1, split2))
  
    def _assertNumSamples(self, builder):
      for split_name, expected_num_examples in self.SPLITS.items():
        self.assertEqual(
            builder.info.splits[split_name].num_examples,
            expected_num_examples,
        )
      self.assertEqual(
          builder.info.splits.total_num_examples,
          sum(self.SPLITS.values()),
      )
  
    def _test_description_builder_config(self, builder):
      # Do not test on external datasets for backward compatibility
      # TODO(tfds): Update once datasets are migrated
      if "tensorflow_datasets" not in self.DATASET_CLASS.__module__:
        return
  
      # If configs specified, ensure they are all valid
      if builder.builder_config and builder.builder_config.description:
        err_msg = textwrap.dedent(
            """\
            The BuilderConfig description should be a one-line description of
            the config.
            It shouldn't be the same as `builder.info.description` to avoid
            redundancy. Both `config.description` and `builder.info.description`
            will be displayed in the catalog.
            """
        )
        ratio = difflib.SequenceMatcher(
            None,
            builder.builder_config.description,
            builder.info.description,
        ).ratio()
        if ratio > 0.9:
          raise AssertionError(err_msg)
  
  
  def checksum(example):
    """Computes the md5 for a given example."""
  
    def _bytes_flatten(flat_str, element):
      """Recursively flatten an element to its byte representation."""
      if isinstance(element, numbers.Number):
        # In python3, bytes(-3) is not allowed (or large numbers),
        # so convert to str to avoid problems.
        element = str(element)
      if isinstance(element, dict):
        for k, v in sorted(element.items()):
          flat_str.append(k)
          _bytes_flatten(flat_str, v)
      elif isinstance(element, str):
        if hasattr(element, "decode"):
          # Python2 considers bytes to be str, but are almost always latin-1
          # encoded bytes here. Extra step needed to avoid DecodeError.
          element = element.decode("latin-1")
        flat_str.append(element)
      elif isinstance(element,
                      (tf.RaggedTensor, tf.compat.v1.ragged.RaggedTensorValue)):
        flat_str.append(str(element.to_list()))
      elif isinstance(element, (np.ndarray, np.generic)):
        # tf.Tensor() returns np.array of dtype object, which don't work
        # with x.to_bytes(). So instead convert numpy into list.
        if element.dtype.type is np.object_:
          flat_str.append(str(tuple(element.shape)))
          flat_str.append(str(list(element.ravel())))
        else:
          flat_str.append(element.tobytes())
      else:
        flat_str.append(bytes(element))
      return flat_str
  
    flat_str = _bytes_flatten([], example)
    flat_bytes = [
        s.encode("utf-8") if not isinstance(s, bytes) else s
        for s in flat_str
    ]
    flat_bytes = b"".join(flat_bytes)
  
    hash_ = hashlib.md5()
    hash_.update(flat_bytes)
    return hash_.hexdigest()
  
  
  def compare_shapes_and_types(tensor_info, output_types, output_shapes):
    """Compare shapes and types between TensorInfo and Dataset types/shapes."""
    for feature_name, feature_info in tensor_info.items():
      if isinstance(feature_info, dict):
        compare_shapes_and_types(feature_info, output_types[feature_name],
                                 output_shapes[feature_name])
      else:
        expected_type = feature_info.dtype
        output_type = output_types[feature_name]
        if expected_type != output_type:
          raise TypeError("Feature %s has type %s but expected %s" %
                          (feature_name, output_type, expected_type))
  
        expected_shape = feature_info.shape
        output_shape = output_shapes[feature_name]
        utils.assert_shape_match(expected_shape, output_shape)
  