


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      UpperCAmelCase_ ( UpperCamelCase   ):

   '''simple docstring'''




   def _snake_case   (		self						,    __A						,    __A=False			):






     """simple docstring"""





     lowerCamelCase   ,     lowerCamelCase   ,     lowerCamelCase			: str   =    (
         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(__A			)						,    "wb"			).close()
     lowerCamelCase			: List[Any]   =    pa.Table.from_pydict({"filename": [Path(__A			).name] * 100}			)
     if take == -1:
       lowerCamelCase			: int   =    len(__A			) - skip
     if skip is not None and take is not None:
       lowerCamelCase			: Optional[int]   =    pa_table.slice(__A						,    __A			)
     return pa_table



class      UpperCAmelCase_ ( UpperCamelCase   ):

   '''simple docstring'''




   def _snake_case   (		self			):






     """simple docstring"""





     lowerCamelCase			: Optional[Any]   =    "my_name"
     lowerCamelCase			: Optional[int]   =    SplitInfo(name="train"						,    num_examples=100			)
     lowerCamelCase			: Dict   =    SplitInfo(name="test"						,    num_examples=100			)
     lowerCamelCase			: Optional[int]   =    [train_info, test_info]
     lowerCamelCase			: Optional[Any]   =    SplitDict()
     split_dict.add(__A			)
     split_dict.add(__A			)
     lowerCamelCase			: Tuple   =    DatasetInfo(splits=__A			)

     with tempfile.TemporaryDirectory() as tmp_dir:
       lowerCamelCase			: Union[str, Any]   =    ReaderTest(__A						,    __A			)

       lowerCamelCase			: Tuple   =    "test[:33%]"
       lowerCamelCase			: str   =    Dataset(**reader.read(__A						,    __A						,    __A			)			)
       self.assertEqual(dset["filename"][0]						,    F"""{name}-test"""			)
       self.assertEqual(dset.num_rows						,    33			)
       self.assertEqual(dset.num_columns						,    1			)

       lowerCamelCase			: int   =    ["train", "test[:33%]"]
       lowerCamelCase			: List[Any]   =    [Split.TRAIN, ReadInstruction.from_spec("test[:33%]"			)]
       for instructions in [instructionsa, instructionsa]:
         lowerCamelCase			: List[Any]   =    [reader.read(__A						,    __A						,    __A			) for instr in instructions]
         lowerCamelCase   ,     lowerCamelCase			: Optional[Any]   =    (Dataset(**__A			) 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						,    __A			)
         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						,    __A			)
         self.assertEqual(str(test_dset.split			)						,    "test[:33%]"			)
         del train_dset, test_dset






   def _snake_case   (		self			):






     """simple docstring"""





     lowerCamelCase			: List[str]   =    "my_name"
     lowerCamelCase			: str   =    SplitInfo(name="train"						,    num_examples=1000						,    shard_lengths=[100] * 10			)
     lowerCamelCase			: Dict   =    [train_info]
     lowerCamelCase			: Optional[int]   =    SplitDict()
     split_dict.add(__A			)
     lowerCamelCase			: List[Any]   =    DatasetInfo(splits=__A			)

     with tempfile.TemporaryDirectory() as tmp_dir:
       lowerCamelCase			: int   =    ReaderTest(__A						,    __A			)

       lowerCamelCase			: List[Any]   =    "train[:33%]"
       lowerCamelCase			: List[Any]   =    Dataset(**reader.read(__A						,    __A						,    __A			)			)
       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 _snake_case   (		self			):






     """simple docstring"""





     lowerCamelCase			: List[str]   =    SplitInfo(name="train"						,    num_examples=100			)
     lowerCamelCase			: Dict   =    SplitInfo(name="test"						,    num_examples=100			)
     lowerCamelCase			: Dict   =    SplitDict()
     split_dict.add(__A			)
     split_dict.add(__A			)
     lowerCamelCase			: Optional[Any]   =    DatasetInfo(splits=__A			)

     with tempfile.TemporaryDirectory() as tmp_dir:
       lowerCamelCase			: List[str]   =    ReaderTest(__A						,    __A			)

       lowerCamelCase			: List[str]   =    [
           {"filename": os.path.join(__A						,    "train"			)},
           {"filename": os.path.join(__A						,    "test"			), "skip": 10, "take": 10},
       ]
       lowerCamelCase			: Optional[int]   =    Dataset(**reader.read_files(__A						,    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 		lowercase_(  SCREAMING_SNAKE_CASE_				,						SCREAMING_SNAKE_CASE_				,						SCREAMING_SNAKE_CASE_     ):


  '''simple docstring'''

  lowerCamelCase			: Tuple   =    arrow_file
  with assert_arrow_memory_increases() if in_memory else assert_arrow_memory_doesnt_increase():
    lowerCamelCase			: List[str]   =    ArrowReader.read_table(SCREAMING_SNAKE_CASE_				,						in_memory=SCREAMING_SNAKE_CASE_     )
  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 		lowercase_(  SCREAMING_SNAKE_CASE_				,						SCREAMING_SNAKE_CASE_				,						SCREAMING_SNAKE_CASE_     ):


  '''simple docstring'''

  lowerCamelCase			: Tuple   =    arrow_file
  lowerCamelCase			: int   =    ArrowReader(""				,						SCREAMING_SNAKE_CASE_     )
  with assert_arrow_memory_increases() if in_memory else assert_arrow_memory_doesnt_increase():
    lowerCamelCase			: Dict   =    reader.read_files([{"filename": filename}]				,						in_memory=SCREAMING_SNAKE_CASE_     )
  assert dataset_kwargs.keys() == {"arrow_table", "info", "split"}
  lowerCamelCase			: List[str]   =    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 		lowercase_(  ):


  '''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%]"

  lowerCamelCase			: int   =    "train+test"
  assert ReadInstruction.from_spec(SCREAMING_SNAKE_CASE_     ).to_spec() == spec_train_test

  lowerCamelCase			: str   =    "train[2:10]"
  assert ReadInstruction.from_spec(SCREAMING_SNAKE_CASE_     ).to_spec() == spec_train_abs

  lowerCamelCase			: Tuple   =    "train[15%:-20%]"
  assert ReadInstruction.from_spec(SCREAMING_SNAKE_CASE_     ).to_spec() == spec_train_pct

  lowerCamelCase			: List[Any]   =    "train[:10%](closest)"
  assert ReadInstruction.from_spec(SCREAMING_SNAKE_CASE_     ).to_spec() == "train[:10%]"

  lowerCamelCase			: Dict   =    "train[:10%](pct1_dropremainder)"
  assert ReadInstruction.from_spec(SCREAMING_SNAKE_CASE_     ).to_spec() == spec_train_pct_rounding

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






def 		lowercase_(  ):


  '''simple docstring'''

  lowerCamelCase			: Tuple   =    "dummy"
  lowerCamelCase			: str   =    [SplitInfo(name="train"				,						num_examples=100     )]
  lowerCamelCase			: Optional[Any]   =    "train[:33%]"
  lowerCamelCase			: Union[str, Any]   =    "arrow"
  lowerCamelCase			: int   =    "prefix"

  lowerCamelCase			: Union[str, Any]   =    make_file_instructions(SCREAMING_SNAKE_CASE_				,						SCREAMING_SNAKE_CASE_				,						SCREAMING_SNAKE_CASE_				,						SCREAMING_SNAKE_CASE_				,						SCREAMING_SNAKE_CASE_     )
  assert isinstance(SCREAMING_SNAKE_CASE_				,						SCREAMING_SNAKE_CASE_     )
  assert file_instructions.num_examples == 33
  assert file_instructions.file_instructions == [
      {"filename": os.path.join(SCREAMING_SNAKE_CASE_				,						f"""{name}-train.arrow"""     ), "skip": 0, "take": 33}
  ]

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







@pytest.mark.parametrize("name, expected_exception"				,						[(None, TypeError), ("", ValueError)]     )
def 		lowercase_(  SCREAMING_SNAKE_CASE_				,						SCREAMING_SNAKE_CASE_     ):


  '''simple docstring'''

  lowerCamelCase			: Tuple   =    [SplitInfo(name="train"				,						num_examples=100     )]
  lowerCamelCase			: Optional[int]   =    "train"
  lowerCamelCase			: Dict   =    "arrow"
  lowerCamelCase			: Optional[Any]   =    "prefix_path"
  with pytest.raises(SCREAMING_SNAKE_CASE_     ):
    lowerCamelCase			: Optional[Any]   =    make_file_instructions(SCREAMING_SNAKE_CASE_				,						SCREAMING_SNAKE_CASE_				,						SCREAMING_SNAKE_CASE_				,						SCREAMING_SNAKE_CASE_				,						SCREAMING_SNAKE_CASE_     )







from typing import Optional, Union

from ...file_utils import TensorType
from ...processing_utils import ProcessorMixin
from ...tokenization_utils_base import BatchEncoding, PaddingStrategy, TruncationStrategy
class      UpperCAmelCase_ ( UpperCamelCase   ):

   '''simple docstring'''




   __A			:				Union[str, Any]							=		"MarkupLMFeatureExtractor"
   __A			:				List[str]							=		("MarkupLMTokenizer", "MarkupLMTokenizerFast")
   __A			:				str							=		True






   def __call__(		self						,    __A=None						,    __A=None						,    __A=None						,    __A=None						,    __A=None						,    __A = True						,    __A = False						,    __A = None						,    __A = None						,    __A = 0						,    __A = None						,    __A = None						,    __A = None						,    __A = False						,    __A = False						,    __A = False						,    __A = False						,    __A = True						,    __A = None						,    **__A						,    ):






     """simple docstring"""





     if self.parse_html:
       if html_strings is None:
         raise ValueError("Make sure to pass HTML strings in case `parse_html` is set to `True`"			)

       if nodes is not None or xpaths is not None or node_labels is not None:
         raise ValueError(
             "Please don't pass nodes, xpaths nor node labels in case `parse_html` is set to `True`"			)

       lowerCamelCase			: Any   =    self.feature_extractor(__A			)
       lowerCamelCase			: Optional[int]   =    features["nodes"]
       lowerCamelCase			: Any   =    features["xpaths"]
     else:
       if html_strings is not None:
         raise ValueError("You have passed HTML strings but `parse_html` is set to `False`."			)
       if nodes is None or xpaths is None:
         raise ValueError("Make sure to pass nodes and xpaths in case `parse_html` is set to `False`"			)

        # # second, apply the tokenizer
     if questions is not None and self.parse_html:
       if isinstance(__A						,    __A			):
         lowerCamelCase			: Tuple   =    [questions]  # add batch dimension (as the feature extractor always adds a batch dimension)

     lowerCamelCase			: Tuple   =    self.tokenizer(
         text=questions if questions is not None else nodes						,    text_pair=nodes if questions is not None else None						,    xpaths=__A						,    node_labels=__A						,    add_special_tokens=__A						,    padding=__A						,    truncation=__A						,    max_length=__A						,    stride=__A						,    pad_to_multiple_of=__A						,    return_token_type_ids=__A						,    return_attention_mask=__A						,    return_overflowing_tokens=__A						,    return_special_tokens_mask=__A						,    return_offsets_mapping=__A						,    return_length=__A						,    verbose=__A						,    return_tensors=__A						,    **__A						,    )

     return encoded_inputs






   def _snake_case   (		self						,    *__A						,    **__A			):






     """simple docstring"""





     return self.tokenizer.batch_decode(*__A						,    **__A			)






   def _snake_case   (		self						,    *__A						,    **__A			):






     """simple docstring"""





     return self.tokenizer.decode(*__A						,    **__A			)






   @property
   def _snake_case   (		self			):






     """simple docstring"""





     lowerCamelCase			: Dict   =    self.tokenizer.model_input_names
     return tokenizer_input_names



