






'''simple docstring'''







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      __SCREAMING_SNAKE_CASE					(__A					):


  """simple docstring"""


  def       _a					( self      ,  UpperCamelCase__      ,  UpperCamelCase__=False  ):






         """simple docstring"""


         a_    ,  a_    ,  a_	    = (
             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()
         a_	    = pa.Table.from_pydict({'filename': [Path(UpperCamelCase__  ).name] * 100}  )
         if take == -1:
                a_	    = len(UpperCamelCase__  ) - skip
         if skip is not None and take is not None:
                a_	    = pa_table.slice(UpperCamelCase__      ,  UpperCamelCase__  )
         return pa_table





class      __SCREAMING_SNAKE_CASE					(__A					):


  """simple docstring"""


  def       _a					( self  ):






         """simple docstring"""


         a_	    = 'my_name'
         a_	    = SplitInfo(name='train'      ,  num_examples=100  )
         a_	    = SplitInfo(name='test'      ,  num_examples=100  )
         a_	    = [train_info, test_info]
         a_	    = SplitDict()
         split_dict.add(UpperCamelCase__  )
         split_dict.add(UpperCamelCase__  )
         a_	    = DatasetInfo(splits=UpperCamelCase__  )

         with tempfile.TemporaryDirectory() as tmp_dir:
                a_	    = ReaderTest(UpperCamelCase__      ,  UpperCamelCase__  )

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

                a_	    = ['train', 'test[:33%]']
                a_	    = [Split.TRAIN, ReadInstruction.from_spec('test[:33%]'  )]
                for instructions in [instructionsa, instructionsa]:
                       a_	    = [reader.read(UpperCamelCase__      ,  UpperCamelCase__      ,  UpperCamelCase__  ) for instr in instructions]
                       a_    ,  a_	    = (Dataset(**UpperCamelCase__  ) for dataset_kwargs in datasets_kwargs)
                       self.assertEqual(train_dset['filename'][0]      ,  f'{name}-train'  )
                       self.assertEqual(train_dset.num_rows      ,  100  )
                       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      ,  33  )
                       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       _a					( self  ):






         """simple docstring"""


         a_	    = 'my_name'
         a_	    = SplitInfo(name='train'      ,  num_examples=1_000      ,  shard_lengths=[100] * 10  )
         a_	    = [train_info]
         a_	    = SplitDict()
         split_dict.add(UpperCamelCase__  )
         a_	    = DatasetInfo(splits=UpperCamelCase__  )

         with tempfile.TemporaryDirectory() as tmp_dir:
                a_	    = ReaderTest(UpperCamelCase__      ,  UpperCamelCase__  )

                a_	    = 'train[:33%]'
                a_	    = 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      ,  330  )
                self.assertEqual(dset.num_columns      ,  1  )







  def       _a					( self  ):






         """simple docstring"""


         a_	    = SplitInfo(name='train'      ,  num_examples=100  )
         a_	    = SplitInfo(name='test'      ,  num_examples=100  )
         a_	    = SplitDict()
         split_dict.add(UpperCamelCase__  )
         split_dict.add(UpperCamelCase__  )
         a_	    = DatasetInfo(splits=UpperCamelCase__  )

         with tempfile.TemporaryDirectory() as tmp_dir:
                a_	    = ReaderTest(UpperCamelCase__      ,  UpperCamelCase__  )

                a_	    = [
                    {'filename': os.path.join(UpperCamelCase__      ,  'train'  )},
                    {'filename': os.path.join(UpperCamelCase__      ,  'test'  ), 'skip': 10, 'take': 10},
                ]
                a_	    = Dataset(**reader.read_files(UpperCamelCase__      ,  original_instructions='train+test[10:20]'  )  )
                self.assertEqual(dset.num_rows      ,  110  )
                self.assertEqual(dset.num_columns      ,  1  )
                del dset





@pytest.mark.parametrize('in_memory'      ,					[False, True]      )
def 			__UpperCamelCase				(			lowercase_					:     List[Any]      ,					lowercase_					:     Dict      ,					lowercase_					:     str      ):
       """simple docstring"""




       a_	    = arrow_file
       with assert_arrow_memory_increases() if in_memory else assert_arrow_memory_doesnt_increase():
              a_	    = ArrowReader.read_table(lowercase_      ,					in_memory=lowercase_      )
       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 			__UpperCamelCase				(			lowercase_					:     int      ,					lowercase_					:     Tuple      ,					lowercase_					:     Optional[Any]      ):
       """simple docstring"""




       a_	    = arrow_file
       a_	    = ArrowReader(''      ,					lowercase_      )
       with assert_arrow_memory_increases() if in_memory else assert_arrow_memory_doesnt_increase():
              a_	    = reader.read_files([{'filename': filename}]      ,					in_memory=lowercase_      )
       assert dataset_kwargs.keys() == {"arrow_table", "info", "split"}
       a_	    = 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 			__UpperCamelCase				(			):
       """simple docstring"""




       assert ReadInstruction('train'      ,					to=10      ,					unit='abs'      ).to_spec() == "train[:10]"
       assert ReadInstruction('train'      ,					from_=-80      ,					to=10      ,					unit='%'      ).to_spec() == "train[-80%:10%]"

       a_	    = 'train+test'
       assert ReadInstruction.from_spec(lowercase_      ).to_spec() == spec_train_test

       a_	    = 'train[2:10]'
       assert ReadInstruction.from_spec(lowercase_      ).to_spec() == spec_train_abs

       a_	    = 'train[15%:-20%]'
       assert ReadInstruction.from_spec(lowercase_      ).to_spec() == spec_train_pct

       a_	    = 'train[:10%](closest)'
       assert ReadInstruction.from_spec(lowercase_      ).to_spec() == "train[:10%]"

       a_	    = 'train[:10%](pct1_dropremainder)'
       assert ReadInstruction.from_spec(lowercase_      ).to_spec() == spec_train_pct_rounding

       a_	    = 'train[:10%](pct1_dropremainder)+test[-10%:](pct1_dropremainder)'
       assert ReadInstruction.from_spec(lowercase_      ).to_spec() == spec_train_test_pct_rounding


def 			__UpperCamelCase				(			):
       """simple docstring"""




       a_	    = 'dummy'
       a_	    = [SplitInfo(name='train'      ,					num_examples=100      )]
       a_	    = 'train[:33%]'
       a_	    = 'arrow'
       a_	    = 'prefix'

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

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


@pytest.mark.parametrize('name, expected_exception'      ,					[(None, TypeError), ('', ValueError)]      )
def 			__UpperCamelCase				(			lowercase_					:     Dict      ,					lowercase_					:     List[Any]      ):
       """simple docstring"""




       a_	    = [SplitInfo(name='train'      ,					num_examples=100      )]
       a_	    = 'train'
       a_	    = 'arrow'
       a_	    = 'prefix_path'
       with pytest.raises(lowercase_      ):
              a_	    = make_file_instructions(lowercase_      ,					lowercase_      ,					lowercase_      ,					lowercase_      ,					lowercase_      )












'''simple docstring'''







import re


def 			__UpperCamelCase				(			lowercase_					:     int      ):
       """simple docstring"""




       def docstring_decorator(lowercase_					:     int      ):
              a_	    = fn.__doc__
              a_	    = func_doc.split('\n'      )
              a_	    = 0
              while i < len(lowercase_      ) and re.search(R'^\s*Examples?:\s*$'      ,					lines[i]      ) is None:
                     i += 1
              if i < len(lowercase_      ):
                     a_	    = example_docstring
                     a_	    = '\n'.join(lowercase_      )
              else:
                     raise ValueError(
                         F'The function {fn} should have an empty \'Examples:\' in its docstring as placeholder, '
                         F'current docstring is:\n{func_doc}'      )
              a_	    = func_doc
              return fn

       return docstring_decorator




