
import os
import tempfile
from pathlib import Path
from unittest import TestCase

import pyarrow as pa
import pytest

from datasets.arrow_dataset import Dataset
from datasets.arrow_reader import ArrowReader, BaseReader, FileInstructions, ReadInstruction, make_file_instructions
from datasets.info import DatasetInfo
from datasets.splits import NamedSplit, Split, SplitDict, SplitInfo

from .utils import assert_arrow_memory_doesnt_increase, assert_arrow_memory_increases

class 		lowercase_       (       lowercase_				):

 def lowerCamelCase_				(      self     ,	__UpperCamelCase     ,	__UpperCamelCase=False     ):




        """simple docstring"""


        UpperCamelCase_       ,				UpperCamelCase_       ,				UpperCamelCase_							       =       (
            filename_skip_take["""filename"""],
            filename_skip_take["""skip"""] if """skip""" in filename_skip_take else None,
            filename_skip_take["""take"""] if """take""" in filename_skip_take else None,
        )
        open(os.path.join(__UpperCamelCase     )     ,	"""wb"""     ).close()
        UpperCamelCase_							       =       pa.Table.from_pydict({"""filename""": [Path(__UpperCamelCase     ).name] * 1_0_0}     )
        if take == -1:
               UpperCamelCase_							       =       len(__UpperCamelCase     ) - skip
        if skip is not None and take is not None:
               UpperCamelCase_							       =       pa_table.slice(__UpperCamelCase     ,	__UpperCamelCase     )
        return pa_table



class 		lowercase_       (       lowercase_				):

 def lowerCamelCase_				(      self     ):




        """simple docstring"""


        UpperCamelCase_							       =       """my_name"""
        UpperCamelCase_							       =       SplitInfo(name="""train"""     ,	num_examples=1_0_0     )
        UpperCamelCase_							       =       SplitInfo(name="""test"""     ,	num_examples=1_0_0     )
        UpperCamelCase_							       =       [train_info, test_info]
        UpperCamelCase_							       =       SplitDict()
        split_dict.add(__UpperCamelCase     )
        split_dict.add(__UpperCamelCase     )
        UpperCamelCase_							       =       DatasetInfo(splits=__UpperCamelCase     )

        with tempfile.TemporaryDirectory() as tmp_dir:
               UpperCamelCase_							       =       ReaderTest(__UpperCamelCase     ,	__UpperCamelCase     )

               UpperCamelCase_							       =       """test[:33%]"""
               UpperCamelCase_							       =       Dataset(**reader.read(__UpperCamelCase     ,	__UpperCamelCase     ,	__UpperCamelCase     )     )
               self.assertEqual(dset["""filename"""][0]     ,	f'''{name}-test'''     )
               self.assertEqual(dset.num_rows     ,	3_3     )
               self.assertEqual(dset.num_columns     ,	1     )

               UpperCamelCase_							       =       ["""train""", """test[:33%]"""]
               UpperCamelCase_							       =       [Split.TRAIN, ReadInstruction.from_spec("""test[:33%]"""     )]
               for instructions in [instructionsa, instructionsa]:
                      UpperCamelCase_							       =       [reader.read(__UpperCamelCase     ,	__UpperCamelCase     ,	__UpperCamelCase     ) for instr in instructions]
                      UpperCamelCase_       ,				UpperCamelCase_							       =       (Dataset(**__UpperCamelCase     ) for dataset_kwargs in datasets_kwargs)
                      self.assertEqual(train_dset["""filename"""][0]     ,	f'''{name}-train'''     )
                      self.assertEqual(train_dset.num_rows     ,	1_0_0     )
                      self.assertEqual(train_dset.num_columns     ,	1     )
                      self.assertIsInstance(train_dset.split     ,	__UpperCamelCase     )
                      self.assertEqual(str(train_dset.split     )     ,	"""train"""     )
                      self.assertEqual(test_dset["""filename"""][0]     ,	f'''{name}-test'''     )
                      self.assertEqual(test_dset.num_rows     ,	3_3     )
                      self.assertEqual(test_dset.num_columns     ,	1     )
                      self.assertIsInstance(test_dset.split     ,	__UpperCamelCase     )
                      self.assertEqual(str(test_dset.split     )     ,	"""test[:33%]"""     )
                      del train_dset, test_dset

 def lowerCamelCase_				(      self     ):




        """simple docstring"""


        UpperCamelCase_							       =       """my_name"""
        UpperCamelCase_							       =       SplitInfo(name="""train"""     ,	num_examples=1_0_0_0     ,	shard_lengths=[1_0_0] * 1_0     )
        UpperCamelCase_							       =       [train_info]
        UpperCamelCase_							       =       SplitDict()
        split_dict.add(__UpperCamelCase     )
        UpperCamelCase_							       =       DatasetInfo(splits=__UpperCamelCase     )

        with tempfile.TemporaryDirectory() as tmp_dir:
               UpperCamelCase_							       =       ReaderTest(__UpperCamelCase     ,	__UpperCamelCase     )

               UpperCamelCase_							       =       """train[:33%]"""
               UpperCamelCase_							       =       Dataset(**reader.read(__UpperCamelCase     ,	__UpperCamelCase     ,	__UpperCamelCase     )     )
               self.assertEqual(dset["""filename"""][0]     ,	f'''{name}-train-00000-of-00010'''     )
               self.assertEqual(dset["""filename"""][-1]     ,	f'''{name}-train-00003-of-00010'''     )
               self.assertEqual(dset.num_rows     ,	3_3_0     )
               self.assertEqual(dset.num_columns     ,	1     )

 def lowerCamelCase_				(      self     ):




        """simple docstring"""


        UpperCamelCase_							       =       SplitInfo(name="""train"""     ,	num_examples=1_0_0     )
        UpperCamelCase_							       =       SplitInfo(name="""test"""     ,	num_examples=1_0_0     )
        UpperCamelCase_							       =       SplitDict()
        split_dict.add(__UpperCamelCase     )
        split_dict.add(__UpperCamelCase     )
        UpperCamelCase_							       =       DatasetInfo(splits=__UpperCamelCase     )

        with tempfile.TemporaryDirectory() as tmp_dir:
               UpperCamelCase_							       =       ReaderTest(__UpperCamelCase     ,	__UpperCamelCase     )

               UpperCamelCase_							       =       [
                   {"""filename""": os.path.join(__UpperCamelCase     ,	"""train"""     )},
                   {"""filename""": os.path.join(__UpperCamelCase     ,	"""test"""     ), """skip""": 1_0, """take""": 1_0},
               ]
               UpperCamelCase_							       =       Dataset(**reader.read_files(__UpperCamelCase     ,	original_instructions="""train+test[10:20]"""     )     )
               self.assertEqual(dset.num_rows     ,	1_1_0     )
               self.assertEqual(dset.num_columns     ,	1     )
               del dset




@pytest.mark.parametrize("""in_memory"""	,      [False, True]     )
def 				lowerCamelCase__				(						a__  :  str	,      a__  :  Tuple	,      a__  :  List[Any]     )		->	Optional[int]:
       UpperCamelCase_							       =       arrow_file
       with assert_arrow_memory_increases() if in_memory else assert_arrow_memory_doesnt_increase():
              UpperCamelCase_							       =       ArrowReader.read_table(a__	,      in_memory=a__     )
       assert table.shape == dataset.data.shape
       assert set(table.column_names     ) == set(dataset.data.column_names     )
       assert dict(table.to_pydict()     ) == dict(dataset.data.to_pydict()     )  # to_pydict returns OrderedDict




@pytest.mark.parametrize("""in_memory"""	,      [False, True]     )
def 				lowerCamelCase__				(						a__  :  Optional[int]	,      a__  :  str	,      a__  :  Optional[int]     )		->	Optional[Any]:
       UpperCamelCase_							       =       arrow_file
       UpperCamelCase_							       =       ArrowReader(""""""	,      a__     )
       with assert_arrow_memory_increases() if in_memory else assert_arrow_memory_doesnt_increase():
              UpperCamelCase_							       =       reader.read_files([{"""filename""": filename}]	,      in_memory=a__     )
       assert dataset_kwargs.keys() == {"arrow_table", "info", "split"}
       UpperCamelCase_							       =       dataset_kwargs["""arrow_table"""]
       assert table.shape == dataset.data.shape
       assert set(table.column_names     ) == set(dataset.data.column_names     )
       assert dict(table.to_pydict()     ) == dict(dataset.data.to_pydict()     )  # to_pydict returns OrderedDict




def 				lowerCamelCase__				(						)		->	Optional[Any]:
       assert ReadInstruction("""train"""	,      to=10	,      unit="""abs"""     ).to_spec() == "train[:10]"
       assert ReadInstruction("""train"""	,      from_=-80	,      to=10	,      unit="""%"""     ).to_spec() == "train[-80%:10%]"

       UpperCamelCase_							       =       """train+test"""
       assert ReadInstruction.from_spec(a__     ).to_spec() == spec_train_test

       UpperCamelCase_							       =       """train[2:10]"""
       assert ReadInstruction.from_spec(a__     ).to_spec() == spec_train_abs

       UpperCamelCase_							       =       """train[15%:-20%]"""
       assert ReadInstruction.from_spec(a__     ).to_spec() == spec_train_pct

       UpperCamelCase_							       =       """train[:10%](closest)"""
       assert ReadInstruction.from_spec(a__     ).to_spec() == "train[:10%]"

       UpperCamelCase_							       =       """train[:10%](pct1_dropremainder)"""
       assert ReadInstruction.from_spec(a__     ).to_spec() == spec_train_pct_rounding

       UpperCamelCase_							       =       """train[:10%](pct1_dropremainder)+test[-10%:](pct1_dropremainder)"""
       assert ReadInstruction.from_spec(a__     ).to_spec() == spec_train_test_pct_rounding




def 				lowerCamelCase__				(						)		->	Optional[Any]:
       UpperCamelCase_							       =       """dummy"""
       UpperCamelCase_							       =       [SplitInfo(name="""train"""	,      num_examples=100     )]
       UpperCamelCase_							       =       """train[:33%]"""
       UpperCamelCase_							       =       """arrow"""
       UpperCamelCase_							       =       """prefix"""

       UpperCamelCase_							       =       make_file_instructions(a__	,      a__	,      a__	,      a__	,      a__     )
       assert isinstance(a__	,      a__     )
       assert file_instructions.num_examples == 33
       assert file_instructions.file_instructions == [
           {"filename": os.path.join(a__	,      f'''{name}-train.arrow'''     ), "skip": 0, "take": 33}
       ]

       UpperCamelCase_							       =       [SplitInfo(name="""train"""	,      num_examples=100	,      shard_lengths=[10] * 10     )]
       UpperCamelCase_							       =       make_file_instructions(a__	,      a__	,      a__	,      a__	,      a__     )
       assert isinstance(a__	,      a__     )
       assert file_instructions.num_examples == 33
       assert file_instructions.file_instructions == [
           {"filename": os.path.join(a__	,      f'''{name}-train-00000-of-00010.arrow'''     ), "skip": 0, "take": -1},
           {"filename": os.path.join(a__	,      f'''{name}-train-00001-of-00010.arrow'''     ), "skip": 0, "take": -1},
           {"filename": os.path.join(a__	,      f'''{name}-train-00002-of-00010.arrow'''     ), "skip": 0, "take": -1},
           {"filename": os.path.join(a__	,      f'''{name}-train-00003-of-00010.arrow'''     ), "skip": 0, "take": 3},
       ]







@pytest.mark.parametrize("""name, expected_exception"""	,      [(None, TypeError), ("""""", ValueError)]     )
def 				lowerCamelCase__				(						a__  :  Union[str, Any]	,      a__  :  Dict     )		->	int:
       UpperCamelCase_							       =       [SplitInfo(name="""train"""	,      num_examples=100     )]
       UpperCamelCase_							       =       """train"""
       UpperCamelCase_							       =       """arrow"""
       UpperCamelCase_							       =       """prefix_path"""
       with pytest.raises(a__     ):
              UpperCamelCase_							       =       make_file_instructions(a__	,      a__	,      a__	,      a__	,      a__     )





import importlib
import inspect
from functools import wraps
from typing import TYPE_CHECKING, Optional, Union

from .download.streaming_download_manager import (
    xbasename,
    xdirname,
    xet_parse,
    xexists,
    xgetsize,
    xglob,
    xgzip_open,
    xisdir,
    xisfile,
    xjoin,
    xlistdir,
    xnumpy_load,
    xopen,
    xpandas_read_csv,
    xpandas_read_excel,
    xPath,
    xrelpath,
    xsio_loadmat,
    xsplit,
    xsplitext,
    xwalk,
    xxml_dom_minidom_parse,
)
from .utils.logging import get_logger
from .utils.patching import patch_submodule
from .utils.py_utils import get_imports


_A	  =					get_logger(__name__)


if TYPE_CHECKING:
      from .builder import DatasetBuilder




def 				lowerCamelCase__				(						a__  :  int	,      a__  :  Optional[Union[str, bool]] = None     )		->	Dict:
       UpperCamelCase_							       =       importlib.import_module(a__     )

       # TODO(QL): always update the module to add subsequent new authentication
       if hasattr(a__	,      """_patched_for_streaming"""     ) and module._patched_for_streaming:
              return

       def wrap_auth(a__  :  int     ):
              @wraps(a__     )
              def wrapper(*a__  :  int	,      **a__  :  List[str]     ):
                     return function(*a__	,      use_auth_token=a__	,      **a__     )

              UpperCamelCase_							       =       """wrap_auth"""
              return wrapper

       # open files in a streaming fashion
       patch_submodule(a__	,      """open"""	,      wrap_auth(a__     )     ).start()
       patch_submodule(a__	,      """os.listdir"""	,      wrap_auth(a__     )     ).start()
       patch_submodule(a__	,      """os.walk"""	,      wrap_auth(a__     )     ).start()
       patch_submodule(a__	,      """glob.glob"""	,      wrap_auth(a__     )     ).start()
       # allow to navigate in remote zip files
       patch_submodule(a__	,      """os.path.join"""	,      a__     ).start()
       patch_submodule(a__	,      """os.path.dirname"""	,      a__     ).start()
       patch_submodule(a__	,      """os.path.basename"""	,      a__     ).start()
       patch_submodule(a__	,      """os.path.relpath"""	,      a__     ).start()
       patch_submodule(a__	,      """os.path.split"""	,      a__     ).start()
       patch_submodule(a__	,      """os.path.splitext"""	,      a__     ).start()
       # allow checks on paths
       patch_submodule(a__	,      """os.path.exists"""	,      wrap_auth(a__     )     ).start()
       patch_submodule(a__	,      """os.path.isdir"""	,      wrap_auth(a__     )     ).start()
       patch_submodule(a__	,      """os.path.isfile"""	,      wrap_auth(a__     )     ).start()
       patch_submodule(a__	,      """os.path.getsize"""	,      wrap_auth(a__     )     ).start()
       patch_submodule(a__	,      """pathlib.Path"""	,      a__     ).start()
       # file readers
       patch_submodule(a__	,      """gzip.open"""	,      wrap_auth(a__     )     ).start()
       patch_submodule(a__	,      """numpy.load"""	,      wrap_auth(a__     )     ).start()
       patch_submodule(a__	,      """pandas.read_csv"""	,      wrap_auth(a__     )	,      attrs=["""__version__"""]     ).start()
       patch_submodule(a__	,      """pandas.read_excel"""	,      wrap_auth(a__     )	,      attrs=["""__version__"""]     ).start()
       patch_submodule(a__	,      """scipy.io.loadmat"""	,      wrap_auth(a__     )	,      attrs=["""__version__"""]     ).start()
       patch_submodule(a__	,      """xml.etree.ElementTree.parse"""	,      wrap_auth(a__     )     ).start()
       patch_submodule(a__	,      """xml.dom.minidom.parse"""	,      wrap_auth(a__     )     ).start()
       UpperCamelCase_							       =       True







def 				lowerCamelCase__				(						a__  :  "DatasetBuilder"     )		->	Tuple:
       extend_module_for_streaming(builder.__module__	,      use_auth_token=builder.use_auth_token     )
       # if needed, we also have to extend additional internal imports (like wmt14 -> wmt_utils)
       if not builder.__module__.startswith("""datasets."""     ):  # check that it's not a packaged builder like csv
              for imports in get_imports(inspect.getfile(builder.__class__     )     ):
                     if imports[0] == "internal":
                            UpperCamelCase_							       =       imports[1]
                            UpperCamelCase_							       =       """.""".join(builder.__module__.split("""."""     )[:-1] + [internal_import_name]     )
                            extend_module_for_streaming(a__	,      use_auth_token=builder.use_auth_token     )

    # builders can inherit from other builders that might use streaming functionality
    # (for example, ImageFolder and AudioFolder inherit from FolderBuilder which implements examples generation)
    # but these parents builders are not patched automatically as they are not instantiated, so we patch them here
       from .builder import DatasetBuilder

       UpperCamelCase_							       =       [
           cls.__module__
           for cls in type(a__     ).__mro__[1:]  # make sure it's not the same module we've already patched
           if issubclass(cls	,      a__     ) and cls.__module__ != DatasetBuilder.__module__
       ]  # check it's not a standard builder from datasets.builder
       for module in parent_builder_modules:
              extend_module_for_streaming(a__	,      use_auth_token=builder.use_auth_token     )



