





import os
from typing import BinaryIO, Optional, Union

import numpy as np
import pyarrow.parquet as pq

from .. import Audio, Dataset, Features, Image, NamedSplit, Value, config
from ..features.features import FeatureType, _visit
from ..formatting import query_table
from ..packaged_modules import _PACKAGED_DATASETS_MODULES
from ..packaged_modules.parquet.parquet import Parquet
from ..utils import logging
from ..utils.typing import NestedDataStructureLike, PathLike
from .abc import AbstractDatasetReader




def       __UpperCamelCase  (_SCREAMING_SNAKE_CASE		)    ->       Optional[int]:
  lowercase__           =  np.inf

  def set_batch_size(_SCREAMING_SNAKE_CASE		) -> None:
    nonlocal batch_size
    if isinstance(_SCREAMING_SNAKE_CASE		,						_SCREAMING_SNAKE_CASE		):
      lowercase__           =  min(_SCREAMING_SNAKE_CASE		,						config.PARQUET_ROW_GROUP_SIZE_FOR_IMAGE_DATASETS		)
    elif isinstance(_SCREAMING_SNAKE_CASE		,						_SCREAMING_SNAKE_CASE		):
      lowercase__           =  min(_SCREAMING_SNAKE_CASE		,						config.PARQUET_ROW_GROUP_SIZE_FOR_AUDIO_DATASETS		)
    elif isinstance(_SCREAMING_SNAKE_CASE		,						_SCREAMING_SNAKE_CASE		) and feature.dtype == "binary":
      lowercase__           =  min(_SCREAMING_SNAKE_CASE		,						config.PARQUET_ROW_GROUP_SIZE_FOR_BINARY_DATASETS		)

  _visit(_SCREAMING_SNAKE_CASE		,						_SCREAMING_SNAKE_CASE		)

  return None if batch_size is np.inf else batch_size
class     SCREAMING_SNAKE_CASE       (UpperCAmelCase				):


   def __init__(   self      :    str      ,					a      :    NestedDataStructureLike[PathLike]      ,					a      :    Optional[NamedSplit] = None      ,					a      :    Optional[Features] = None      ,					a      :    str = None      ,					a      :    bool = False      ,					a      :    bool = False      ,					a      :    Optional[int] = None      ,					**a      :    List[Any]      ,					)->       Optional[Any]:





     """simple docstring"""







     super().__init__(
         a      ,					split=a      ,					features=a      ,					cache_dir=a      ,					keep_in_memory=a      ,					streaming=a      ,					num_proc=a      ,					**a      ,					)
     lowercase__           =  path_or_paths if isinstance(a      ,					a	) else {self.split: path_or_paths}
     lowercase__           =  _PACKAGED_DATASETS_MODULES['parquet'][1]
     lowercase__           =  Parquet(
         cache_dir=a      ,					data_files=a      ,					features=a      ,					hash=a      ,					**a      ,					)


   def 	SCREAMING_SNAKE_CASE_							(   self      :    Union[str, Any]	)->       List[Any]:





     """simple docstring"""







     if self.streaming:
       lowercase__           =  self.builder.as_streaming_dataset(split=self.split	)
     # Build regular (map-style) dataset
     else:
       lowercase__           =  None
       lowercase__           =  None
       lowercase__           =  None
       lowercase__           =  None

       self.builder.download_and_prepare(
           download_config=a      ,					download_mode=a      ,					verification_mode=a      ,					base_path=a      ,					num_proc=self.num_proc      ,					)
       lowercase__           =  self.builder.as_dataset(
           split=self.split      ,					verification_mode=a      ,					in_memory=self.keep_in_memory	)
     return dataset






class     SCREAMING_SNAKE_CASE       :


   def __init__(   self      :    Optional[Any]      ,					a      :    Dataset      ,					a      :    Union[PathLike, BinaryIO]      ,					a      :    Optional[int] = None      ,					**a      :    List[Any]      ,					)->       Any:





     """simple docstring"""







     lowercase__           =  dataset
     lowercase__           =  path_or_buf
     lowercase__           =  batch_size or get_writer_batch_size(dataset.features	)
     lowercase__           =  parquet_writer_kwargs


   def 	SCREAMING_SNAKE_CASE_							(   self      :    str	)->       int:





     """simple docstring"""







     lowercase__           =  self.batch_size if self.batch_size else config.DEFAULT_MAX_BATCH_SIZE

     if isinstance(self.path_or_buf      ,					(str, bytes, os.PathLike)	):
       with open(self.path_or_buf      ,					'wb+'	) as buffer:
         lowercase__           =  self._write(file_obj=a      ,					batch_size=a      ,					**self.parquet_writer_kwargs	)
     else:
       lowercase__           =  self._write(file_obj=self.path_or_buf      ,					batch_size=a      ,					**self.parquet_writer_kwargs	)
     return written


   def 	SCREAMING_SNAKE_CASE_							(   self      :    Optional[Any]      ,					a      :    BinaryIO      ,					a      :    int      ,					**a      :    Optional[Any]	)->       int:





     """simple docstring"""







     lowercase__           =  0
     lowercase__           =  parquet_writer_kwargs.pop('path_or_buf'      ,					a	)
     lowercase__           =  self.dataset.features.arrow_schema

     lowercase__           =  pq.ParquetWriter(a      ,					schema=a      ,					**a	)

     for offset in logging.tqdm(
         range(0      ,					len(self.dataset	)      ,					a	)      ,					unit='ba'      ,					disable=not logging.is_progress_bar_enabled()      ,					desc='Creating parquet from Arrow format'      ,					):
       lowercase__           =  query_table(
           table=self.dataset._data      ,					key=slice(a      ,					offset + batch_size	)      ,					indices=self.dataset._indices if self.dataset._indices is not None else None      ,					)
       writer.write_table(a	)
       written += batch.nbytes
     writer.close()
     return written







from datasets.utils.patching import _PatchedModuleObj, patch_submodule

from . import _test_patching




def       __UpperCamelCase  ()    ->       str:
  import os as original_os
  from os import path as original_path
  from os import rename as original_rename
  from os.path import dirname as original_dirname
  from os.path import join as original_join

  assert _test_patching.os is original_os
  assert _test_patching.path is original_path
  assert _test_patching.join is original_join

  assert _test_patching.renamed_os is original_os
  assert _test_patching.renamed_path is original_path
  assert _test_patching.renamed_join is original_join

  lowercase__           =  '__test_patch_submodule_mock__'
  with patch_submodule(_test_patching		,						'os.path.join'		,						_SCREAMING_SNAKE_CASE		):
    # Every way to access os.path.join must be patched, and the rest must stay untouched

    # check os.path.join
    assert isinstance(_test_patching.os		,						_PatchedModuleObj		)
    assert isinstance(_test_patching.os.path		,						_PatchedModuleObj		)
    assert _test_patching.os.path.join is mock

    # check path.join
    assert isinstance(_test_patching.path		,						_PatchedModuleObj		)
    assert _test_patching.path.join is mock

    # check join
    assert _test_patching.join is mock

    # check that the other attributes are untouched
    assert _test_patching.os.rename is original_rename
    assert _test_patching.path.dirname is original_dirname
    assert _test_patching.os.path.dirname is original_dirname

    # Even renamed modules or objects must be patched

    # check renamed_os.path.join
    assert isinstance(_test_patching.renamed_os		,						_PatchedModuleObj		)
    assert isinstance(_test_patching.renamed_os.path		,						_PatchedModuleObj		)
    assert _test_patching.renamed_os.path.join is mock

    # check renamed_path.join
    assert isinstance(_test_patching.renamed_path		,						_PatchedModuleObj		)
    assert _test_patching.renamed_path.join is mock

    # check renamed_join
    assert _test_patching.renamed_join is mock

    # check that the other attributes are untouched
    assert _test_patching.renamed_os.rename is original_rename
    assert _test_patching.renamed_path.dirname is original_dirname
    assert _test_patching.renamed_os.path.dirname is original_dirname

  # check that everthing is back to normal when the patch is over

  assert _test_patching.os is original_os
  assert _test_patching.path is original_path
  assert _test_patching.join is original_join

  assert _test_patching.renamed_os is original_os
  assert _test_patching.renamed_path is original_path
  assert _test_patching.renamed_join is original_join




def       __UpperCamelCase  ()    ->       Any:
  assert _test_patching.open is open

  lowercase__           =  '__test_patch_submodule_builtin_mock__'
  # _test_patching has "open" in its globals
  assert _test_patching.open is open
  with patch_submodule(_test_patching		,						'open'		,						_SCREAMING_SNAKE_CASE		):
    assert _test_patching.open is mock

  # check that everthing is back to normal when the patch is over

  assert _test_patching.open is open




def       __UpperCamelCase  ()    ->       List[str]:
  # pandas.read_csv is not present in _test_patching
  lowercase__           =  '__test_patch_submodule_missing_mock__'
  with patch_submodule(_test_patching		,						'pandas.read_csv'		,						_SCREAMING_SNAKE_CASE		):
    pass




def       __UpperCamelCase  ()    ->       List[str]:
  # builtin should always be mocked even if they're not in the globals
  # in case they're loaded at one point
  lowercase__           =  '__test_patch_submodule_missing_builtin_mock__'
  # _test_patching doesn't have "len" in its globals
  assert getattr(_test_patching		,						'len'		,						_SCREAMING_SNAKE_CASE		) is None
  with patch_submodule(_test_patching		,						'len'		,						_SCREAMING_SNAKE_CASE		):
    assert _test_patching.len is mock
  assert _test_patching.len is len




def       __UpperCamelCase  ()    ->       List[str]:
  lowercase__           =  '__test_patch_submodule_start_and_stop_mock__'
  lowercase__           =  patch_submodule(_test_patching		,						'open'		,						_SCREAMING_SNAKE_CASE		)
  assert _test_patching.open is open
  patch.start()
  assert _test_patching.open is mock
  patch.stop()
  assert _test_patching.open is open




def       __UpperCamelCase  ()    ->       Optional[int]:
  from os import rename as original_rename
  from os.path import dirname as original_dirname
  from os.path import join as original_join

  lowercase__           =  '__test_patch_submodule_successive_join__'
  lowercase__           =  '__test_patch_submodule_successive_dirname__'
  lowercase__           =  '__test_patch_submodule_successive_rename__'
  assert _test_patching.os.path.join is original_join
  assert _test_patching.os.path.dirname is original_dirname
  assert _test_patching.os.rename is original_rename

  with patch_submodule(_test_patching		,						'os.path.join'		,						_SCREAMING_SNAKE_CASE		):
    with patch_submodule(_test_patching		,						'os.rename'		,						_SCREAMING_SNAKE_CASE		):
      with patch_submodule(_test_patching		,						'os.path.dirname'		,						_SCREAMING_SNAKE_CASE		):
        assert _test_patching.os.path.join is mock_join
        assert _test_patching.os.path.dirname is mock_dirname
        assert _test_patching.os.rename is mock_rename

    # try another order
  with patch_submodule(_test_patching		,						'os.rename'		,						_SCREAMING_SNAKE_CASE		):
    with patch_submodule(_test_patching		,						'os.path.join'		,						_SCREAMING_SNAKE_CASE		):
      with patch_submodule(_test_patching		,						'os.path.dirname'		,						_SCREAMING_SNAKE_CASE		):
        assert _test_patching.os.path.join is mock_join
        assert _test_patching.os.path.dirname is mock_dirname
        assert _test_patching.os.rename is mock_rename

  assert _test_patching.os.path.join is original_join
  assert _test_patching.os.path.dirname is original_dirname
  assert _test_patching.os.rename is original_rename




def       __UpperCamelCase  ()    ->       Optional[Any]:
  lowercase__           =  '__test_patch_submodule_doesnt_exist_mock__'
  with patch_submodule(_test_patching		,						'__module_that_doesn_exist__.__attribute_that_doesn_exist__'		,						_SCREAMING_SNAKE_CASE		):
    pass
  with patch_submodule(_test_patching		,						'os.__attribute_that_doesn_exist__'		,						_SCREAMING_SNAKE_CASE		):
    pass


