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.
  
  """Util to generate the dataset documentation content.
  
  Used by tensorflow_datasets/scripts/documentation/build_catalog.py
  
  """
  
  import collections
  from concurrent import futures
  import functools
  from typing import Any, Dict, Iterator, List, Optional, Type
  
  import dataclasses
  
  import tensorflow_datasets as tfds
  from tensorflow_datasets.scripts.documentation import dataset_markdown_builder
  from tensorflow_datasets.scripts.documentation import doc_utils
  import tqdm
  
  _WORKER_COUNT_DATASETS = 50
  _WORKER_COUNT_CONFIGS = 20
  
  # WmtTranslate: The raw wmt can only be instantiated with the config kwargs
  _BUILDER_BLACKLIST = [
      'wmt_translate',
  ]
  
  
  @dataclasses.dataclass(eq=False, frozen=True)
  class BuilderToDocument:
    """Structure containing metadata."""
    section: str
    namespace: Optional[str]
    builder: tfds.core.DatasetBuilder
    config_builders: List[tfds.core.DatasetBuilder]
  
  
  @dataclasses.dataclass(eq=False, frozen=True)
  class BuilderDocumentation:
    """Documentation output of a single builder.
  
    Attributes:
      name: Documentation page name (e.g. `mnist`, `kaggle:mnist`)
      filestem: Documentation page name without suffix
        (e.g. `mnist`, `kaggle_mnist`)
      content: Documentation content
      section: Documentation section (e.g `text`, `image`,...)
      namespace: Dataset namespace
      is_manual: Whether the dataset require manual download
      is_nightly: Whether the dataset was recently added in `tfds-nightly`
    """
    name: str
    filestem: str
    content: str
    section: str
    is_manual: bool
    is_nightly: bool
  
  
  def _load_builder(
      name: str,
  ) -> Optional[BuilderToDocument]:
    """Load the builder to document.
  
    Args:
      name: Builder to load
  
    Returns:
      section: The section in which the builder is documented
      builder: Main builder instance
      config_builders: Additional builders (one of each configs)
    """
    if tfds.core.utils.DatasetName(name).namespace:  # Community dataset
      return _load_builder_from_location(name)
    else:  # Code dataset
      return _load_builder_from_code(name)
  
  
  def _load_builder_from_location(
      name: str,
  ) -> Optional[BuilderToDocument]:
    """Load the builder, config,... to document."""
    namespace = tfds.core.utils.DatasetName(name).namespace
    try:
      builder = tfds.builder(name)
    except tfds.core.DatasetNotFoundError:
      # If tfds.core.DatasetNotFoundError, it might be the default
      # config isn't found. Should try to load a sub-folder (to check ).
      builder = _maybe_load_config(name)
      if not builder:
        return builder
    if builder.builder_config:
      config_builders = _load_all_configs(name, builder)
    else:
      config_builders = []
    return BuilderToDocument(
        section=namespace,
        namespace=namespace,
        builder=builder,
        config_builders=config_builders,
    )
  
  
  def _maybe_load_config(name: str) -> Optional[tfds.core.DatasetBuilder]:
    del name  # unused
    return None
  
  
  def _load_all_configs(
      name: str, builder: tfds.core.DatasetBuilder,
  ) -> List[tfds.core.DatasetBuilder]:
    """Load all builder configs."""
    # `data_dir/name/config/version/` -> `data_dir/name/`
    common_dir = builder.data_path.parent.parent
    filtered_dirs = {builder.builder_config.name, '.config'}
  
    def get_config_builder(path):
      if path.name in filtered_dirs:
        return None  # Default config is already loaded
      try:
        return tfds.builder(f'{name}/{path.name}')
      except tfds.core.DatasetNotFoundError:
        return None
  
    with futures.ThreadPoolExecutor(max_workers=_WORKER_COUNT_CONFIGS) as tpool:
      config_names = sorted(common_dir.iterdir())
      config_builders = [
          config_builder
          for config_builder in tpool.map(get_config_builder, config_names)
          if config_builder is not None  # Filter invalid configs
      ]
  
    return [builder] + config_builders
  
  
  def _load_builder_from_code(
      name: str,
  ) -> BuilderToDocument:
    """Load the builder, config,... to document."""
    builder_cls = tfds.builder_cls(name)
    section = _get_section(builder_cls)
  
    if builder_cls.BUILDER_CONFIGS:  # Builder with configs
  
      def get_config_builder(config) -> tfds.core.DatasetBuilder:
        return tfds.builder(builder_cls.name, config=config)
  
      with futures.ThreadPoolExecutor(max_workers=_WORKER_COUNT_CONFIGS) as tpool:
        config_builders = list(
            tpool.map(get_config_builder, builder_cls.BUILDER_CONFIGS),
        )
      return BuilderToDocument(
          section=section,
          namespace=None,
          builder=config_builders[0],
          config_builders=config_builders,
      )
    else:  # Builder without configs
      return BuilderToDocument(
          section=section,
          namespace=None,
          builder=builder_cls(),  # pytype: disable=not-instantiable
          config_builders=[],
      )
  
  
  def _get_section(builder_cls: Type[tfds.core.DatasetBuilder]) -> str:
    """Returns the section associated with the builder."""
    module_parts = builder_cls.__module__.split('.')
    if module_parts[0] != 'tensorflow_datasets':
      raise AssertionError(f'Unexpected builder {builder_cls}: module')
    _, category, *_ = module_parts  # tfds.<category>.xyz
    return category
  
  
  def _document_single_builder(
      name: str, **kwargs: Any,
  ) -> Optional[BuilderDocumentation]:
    """Doc string for a single builder, with or without configs."""
    with tfds.core.utils.try_reraise(f'Error for `{name}`: '):
      return _document_single_builder_inner(name, **kwargs)
  
  
  def _document_single_builder_inner(
      name: str,
      nightly_doc_util: doc_utils.NightlyDocUtil,
      **kwargs: Any,
  ) -> Optional[BuilderDocumentation]:
    """Doc string for a single builder, with or without configs."""
    tqdm.tqdm.write(f'Document builder {name}...')
    doc_info = _load_builder(name)
    if doc_info is None:
      return None
  
    out_str = dataset_markdown_builder.get_markdown_string(
        namespace=doc_info.namespace,
        builder=doc_info.builder,
        config_builders=doc_info.config_builders,
        nightly_doc_util=nightly_doc_util,
        **kwargs
    )
    is_nightly = bool(
        nightly_doc_util and nightly_doc_util.is_builder_nightly(name)
    )
    return BuilderDocumentation(
        name=name,
        filestem=name.replace(':', '_'),
        content=out_str,
        section=doc_info.section,
        is_manual=bool(doc_info.builder.MANUAL_DOWNLOAD_INSTRUCTIONS),
        is_nightly=is_nightly,
    )
  
  
  def _all_tfds_datasets() -> List[str]:
    """Returns all "official" TFDS dataset names."""
    return sorted([
        name for name in tfds.list_builders(with_community_datasets=True)  # pylint: disable=g-complex-comprehension
        if name not in _BUILDER_BLACKLIST
    ])
  
  
  def iter_documentation_builders(
      datasets: Optional[List[str]] = None,
      *,
      doc_util_paths: doc_utils.DocUtilPaths = None,
  ) -> Iterator[BuilderDocumentation]:
    """Create dataset documentation string for given datasets.
  
    Args:
      datasets: list of datasets for which to create documentation.
                If None, then all available datasets will be used.
      doc_util_paths: Additional path for visualization, nightly info,...
  
    Yields:
      builder_documetation: The documentation information for each builder
    """
    print('Retrieving the list of builders...')
    datasets = datasets or _all_tfds_datasets()
  
    if doc_util_paths.fig_base_path:
      visu_doc_util = doc_utils.VisualizationDocUtil(
          base_path=doc_util_paths.fig_base_path,
          base_url=doc_util_paths.fig_base_url,
      )
    else:
      visu_doc_util = None
  
    if doc_util_paths.df_base_path:
      df_doc_util = doc_utils.DataframeDocUtil(
          base_path=doc_util_paths.df_base_path,
          base_url=doc_util_paths.df_base_url,
      )
    else:
      df_doc_util = None
  
    if doc_util_paths.fig_base_path:
      nightly_doc_util = doc_utils.NightlyDocUtil(
          path=doc_util_paths.nightly_path,
      )
    else:
      nightly_doc_util = None
  
    document_single_builder_fn = functools.partial(
        _document_single_builder,
        visu_doc_util=visu_doc_util,
        df_doc_util=df_doc_util,
        nightly_doc_util=nightly_doc_util,
    )
  
    # Document all builders
    print(f'Document {len(datasets)} builders...')
    with futures.ThreadPoolExecutor(max_workers=_WORKER_COUNT_DATASETS) as tpool:
      tasks = [
          tpool.submit(document_single_builder_fn, name) for name in datasets
      ]
      for future in tqdm.tqdm(futures.as_completed(tasks), total=len(tasks)):
        builder_doc = future.result()
        if builder_doc is None:  # Builder filtered
          continue
        else:
          tqdm.tqdm.write(f'Documentation generated for {builder_doc.name}...')
          yield builder_doc
    print('All builder documentations generated!')
  
  
  def make_category_to_builders_dict() -> Dict[
      str, List[tfds.core.DatasetBuilder]
  ]:
    """Loads all builders with their associated category."""
    datasets = _all_tfds_datasets()
    print(f'Creating the vanilla builders for {len(datasets)} datasets...')
    with futures.ThreadPoolExecutor(max_workers=_WORKER_COUNT_DATASETS) as tpool:
      builders = tpool.map(tfds.builder, datasets)
  
    category_to_builders = collections.defaultdict(list)
    for builder in builders:
      section = _get_section(type(builder))
      category_to_builders[section].append(builder)
    return category_to_builders
  