





import copy
import os
import tempfile
from unittest import TestCase
from unittest.mock import patch

import numpy as np
import pyarrow as pa
import pyarrow.parquet as pq
import pytest

from datasets.arrow_writer import ArrowWriter, OptimizedTypedSequence, ParquetWriter, TypedSequence
from datasets.features import ArrayaD, ClassLabel, Features, Image, Value
from datasets.features.features import ArrayaDExtensionType, cast_to_python_objects
from datasets.keyhash import DuplicatedKeysError, InvalidKeyError

from .utils import require_pil

class 		snake_case_   (							__lowercase						):







	'''simple docstring'''




	def    UpperCAmelCase			(						self							:     List[str]   )       -> Optional[int]:

			'''simple docstring'''




			__lowercase											=     pa.array(TypedSequence([1, 2, 3]   )   )
			self.assertEqual(arr.type					,      pa.intaa()   )




	def    UpperCAmelCase			(						self							:     Tuple   )       -> Optional[Any]:

			'''simple docstring'''




			with self.assertRaises(__lowerCamelCase   ):
					__lowercase											=     pa.array(TypedSequence([1, 2, 3]   )					,      type=pa.intaa()   )




	def    UpperCAmelCase			(						self							:     Any   )       -> Optional[int]:

			'''simple docstring'''




			with self.assertRaises(__lowerCamelCase   ):
					__lowercase											=     pa.array(TypedSequence([1, 2, 3]					,      try_type=Value('bool'   )					,      type=Value('int64'   )   )   )




	def    UpperCAmelCase			(						self							:     Dict   )       -> List[Any]:

			'''simple docstring'''




			__lowercase											=     pa.array(TypedSequence([1, 2, 3]					,      type=Value('int32'   )   )   )
			self.assertEqual(arr.type					,      pa.intaa()   )




	def    UpperCAmelCase			(						self							:     Dict   )       -> str:

			'''simple docstring'''




			with self.assertRaises((TypeError, pa.lib.ArrowInvalid)   ):
					__lowercase											=     pa.array(TypedSequence(['foo', 'bar']					,      type=Value('int64'   )   )   )




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

			'''simple docstring'''




			__lowercase											=     pa.array(TypedSequence([1, 2, 3]					,      try_type=Value('int32'   )   )   )
			self.assertEqual(arr.type					,      pa.intaa()   )




	def    UpperCAmelCase			(						self							:     Dict   )       -> Tuple:

			'''simple docstring'''




			__lowercase											=     pa.array(TypedSequence(['foo', 'bar']					,      try_type=Value('int64'   )   )   )
			self.assertEqual(arr.type					,      pa.string()   )




	def    UpperCAmelCase			(						self							:     Any   )       -> int:

			'''simple docstring'''




			__lowercase											=     pa.array(TypedSequence([[[1, 2, 3]]]					,      type=ArrayaD((1, 3)					,      'int64'   )   )   )
			self.assertEqual(arr.type					,      ArrayaDExtensionType((1, 3)					,      'int64'   )   )




	def    UpperCAmelCase			(						self							:     int   )       -> Tuple:

			'''simple docstring'''




			with self.assertRaises((TypeError, pa.lib.ArrowInvalid)   ):
					__lowercase											=     pa.array(TypedSequence(['foo', 'bar']					,      type=ArrayaD((1, 3)					,      'int64'   )   )   )




	def    UpperCAmelCase			(						self							:     Optional[int]   )       -> Optional[Any]:

			'''simple docstring'''




			__lowercase											=     pa.array(TypedSequence([[[1, 2, 3]]]					,      try_type=ArrayaD((1, 3)					,      'int64'   )   )   )
			self.assertEqual(arr.type					,      ArrayaDExtensionType((1, 3)					,      'int64'   )   )




	def    UpperCAmelCase			(						self							:     Tuple   )       -> str:

			'''simple docstring'''




			__lowercase											=     pa.array(TypedSequence(['foo', 'bar']					,      try_type=ArrayaD((1, 3)					,      'int64'   )   )   )
			self.assertEqual(arr.type					,      pa.string()   )






	@require_pil
	def    UpperCAmelCase			(						self							:     Dict   )       -> Dict:

			'''simple docstring'''




			import PIL.Image

			__lowercase											=     PIL.Image.fromarray(np.arange(10					,      dtype=np.uinta   ).reshape(2					,      5   )   )
			with patch(
			    'datasets.arrow_writer.cast_to_python_objects'					,      side_effect=__lowerCamelCase   ) as mock_cast_to_python_objects:
					__lowercase											=     pa.array(TypedSequence([{'path': None, 'bytes': B'image_bytes'}, pil_image]					,      type=Image()   )   )
					__lowercase											=     mock_cast_to_python_objects.call_args_list[-1]
					self.assertIn('optimize_list_casting'					,      __lowerCamelCase   )
					self.assertFalse(kwargs['optimize_list_casting']   )





def  SCREAMING_SNAKE_CASE				(      snake_case						,   snake_case							)		->    Any:
		__lowercase											=     pa.BufferReader(_lowerCamelCase							) if isinstance(_lowerCamelCase						,   pa.Buffer							) else pa.memory_map(_lowerCamelCase							)
		__lowercase											=     pa.ipc.open_stream(_lowerCamelCase							)
		__lowercase											=     f.read_all()
		assert len(pa_table.to_batches()							) == expected_num_chunks
		assert pa_table.to_pydict() == {"col_1": ["foo", "bar"], "col_2": [1, 2]}
		del pa_table



@pytest.mark.parametrize('writer_batch_size'						,   [None, 1, 10]							)
@pytest.mark.parametrize(
    'fields'						,   [None, {'col_1': pa.string(), 'col_2': pa.intaa()}, {'col_1': pa.string(), 'col_2': pa.intaa()}]							)
def  SCREAMING_SNAKE_CASE				(      snake_case						,   snake_case							)		->    Optional[Any]:
		__lowercase											=     pa.BufferOutputStream()
		__lowercase											=     pa.schema(_lowerCamelCase							) if fields else None
		with ArrowWriter(stream=_lowerCamelCase						,   schema=_lowerCamelCase						,   writer_batch_size=_lowerCamelCase							) as writer:
				writer.write({'col_1': 'foo', 'col_2': 1}							)
				writer.write({'col_1': 'bar', 'col_2': 2}							)
				__lowercase											=     writer.finalize()
		assert num_examples == 2
		assert num_bytes > 0
		if not fields:
				__lowercase											=     {"col_1": pa.string(), "col_2": pa.intaa()}
		assert writer._schema == pa.schema(_lowerCamelCase						,   metadata=writer._schema.metadata							)
		_check_output(output.getvalue()						,   expected_num_chunks=num_examples if writer_batch_size == 1 else 1							)



def  SCREAMING_SNAKE_CASE				(      )		->    Any:
		__lowercase											=     pa.BufferOutputStream()
		__lowercase											=     Features({'labels': ClassLabel(names=['neg', 'pos']							)}							)
		with ArrowWriter(stream=_lowerCamelCase						,   features=_lowerCamelCase							) as writer:
				writer.write({'labels': 0}							)
				writer.write({'labels': 1}							)
				__lowercase											=     writer.finalize()
		assert num_examples == 2
		assert num_bytes > 0
		assert writer._schema == features.arrow_schema
		assert writer._schema.metadata == features.arrow_schema.metadata
		__lowercase											=     pa.BufferReader(output.getvalue()							)
		__lowercase											=     pa.ipc.open_stream(_lowerCamelCase							)
		__lowercase											=     f.read_all()
		__lowercase											=     pa_table.schema
		assert pa_table.num_rows == 2
		assert schema == features.arrow_schema
		assert schema.metadata == features.arrow_schema.metadata
		assert features == Features.from_arrow_schema(_lowerCamelCase							)



@pytest.mark.parametrize('writer_batch_size'						,   [None, 1, 10]							)
def  SCREAMING_SNAKE_CASE				(      snake_case							)		->    List[Any]:
		__lowercase											=     pa.BufferOutputStream()
		with ArrowWriter(
		    stream=_lowerCamelCase						,   writer_batch_size=_lowerCamelCase						,   hash_salt='split_name'						,   check_duplicates=_lowerCamelCase						,   ) as writer:
				with pytest.raises(_lowerCamelCase							):
						writer.write({'col_1': 'foo', 'col_2': 1}						,   key=[1, 2]							)
						__lowercase											=     writer.finalize()



@pytest.mark.parametrize('writer_batch_size'						,   [None, 2, 10]							)
def  SCREAMING_SNAKE_CASE				(      snake_case							)		->    Any:
		__lowercase											=     pa.BufferOutputStream()
		with ArrowWriter(
		    stream=_lowerCamelCase						,   writer_batch_size=_lowerCamelCase						,   hash_salt='split_name'						,   check_duplicates=_lowerCamelCase						,   ) as writer:
				with pytest.raises(_lowerCamelCase							):
						writer.write({'col_1': 'foo', 'col_2': 1}						,   key=10							)
						writer.write({'col_1': 'bar', 'col_2': 2}						,   key=10							)
						__lowercase											=     writer.finalize()



@pytest.mark.parametrize('writer_batch_size'						,   [None, 2, 10]							)
def  SCREAMING_SNAKE_CASE				(      snake_case							)		->    Optional[Any]:
		__lowercase											=     pa.BufferOutputStream()
		with ArrowWriter(
		    stream=_lowerCamelCase						,   writer_batch_size=_lowerCamelCase						,   hash_salt='split_name'						,   check_duplicates=_lowerCamelCase						,   ) as writer:
				writer.write({'col_1': 'foo', 'col_2': 1}						,   key=1							)
				writer.write({'col_1': 'bar', 'col_2': 2}						,   key=2							)
				__lowercase											=     writer.finalize()
		assert num_examples == 2
		assert num_bytes > 0
		_check_output(output.getvalue()						,   expected_num_chunks=num_examples if writer_batch_size == 1 else 1							)



@pytest.mark.parametrize('writer_batch_size'						,   [None, 1, 10]							)
@pytest.mark.parametrize(
    'fields'						,   [None, {'col_1': pa.string(), 'col_2': pa.intaa()}, {'col_1': pa.string(), 'col_2': pa.intaa()}]							)
def  SCREAMING_SNAKE_CASE				(      snake_case						,   snake_case							)		->    str:
		__lowercase											=     pa.BufferOutputStream()
		__lowercase											=     pa.schema(_lowerCamelCase							) if fields else None
		with ArrowWriter(stream=_lowerCamelCase						,   schema=_lowerCamelCase						,   writer_batch_size=_lowerCamelCase							) as writer:
				writer.write_batch({'col_1': ['foo', 'bar'], 'col_2': [1, 2]}							)
				writer.write_batch({'col_1': [], 'col_2': []}							)
				__lowercase											=     writer.finalize()
		assert num_examples == 2
		assert num_bytes > 0
		if not fields:
				__lowercase											=     {"col_1": pa.string(), "col_2": pa.intaa()}
		assert writer._schema == pa.schema(_lowerCamelCase						,   metadata=writer._schema.metadata							)
		_check_output(output.getvalue()						,   expected_num_chunks=num_examples if writer_batch_size == 1 else 1							)



@pytest.mark.parametrize('writer_batch_size'						,   [None, 1, 10]							)
@pytest.mark.parametrize(
    'fields'						,   [None, {'col_1': pa.string(), 'col_2': pa.intaa()}, {'col_1': pa.string(), 'col_2': pa.intaa()}]							)
def  SCREAMING_SNAKE_CASE				(      snake_case						,   snake_case							)		->    Dict:
		__lowercase											=     pa.BufferOutputStream()
		__lowercase											=     pa.schema(_lowerCamelCase							) if fields else None
		with ArrowWriter(stream=_lowerCamelCase						,   schema=_lowerCamelCase						,   writer_batch_size=_lowerCamelCase							) as writer:
				writer.write_table(pa.Table.from_pydict({'col_1': ['foo', 'bar'], 'col_2': [1, 2]}							)							)
				__lowercase											=     writer.finalize()
		assert num_examples == 2
		assert num_bytes > 0
		if not fields:
				__lowercase											=     {"col_1": pa.string(), "col_2": pa.intaa()}
		assert writer._schema == pa.schema(_lowerCamelCase						,   metadata=writer._schema.metadata							)
		_check_output(output.getvalue()						,   expected_num_chunks=num_examples if writer_batch_size == 1 else 1							)



@pytest.mark.parametrize('writer_batch_size'						,   [None, 1, 10]							)
@pytest.mark.parametrize(
    'fields'						,   [None, {'col_1': pa.string(), 'col_2': pa.intaa()}, {'col_1': pa.string(), 'col_2': pa.intaa()}]							)
def  SCREAMING_SNAKE_CASE				(      snake_case						,   snake_case							)		->    Optional[Any]:
		__lowercase											=     pa.BufferOutputStream()
		__lowercase											=     pa.schema(_lowerCamelCase							) if fields else None
		with ArrowWriter(stream=_lowerCamelCase						,   schema=_lowerCamelCase						,   writer_batch_size=_lowerCamelCase							) as writer:
				writer.write_row(pa.Table.from_pydict({'col_1': ['foo'], 'col_2': [1]}							)							)
				writer.write_row(pa.Table.from_pydict({'col_1': ['bar'], 'col_2': [2]}							)							)
				__lowercase											=     writer.finalize()
		assert num_examples == 2
		assert num_bytes > 0
		if not fields:
				__lowercase											=     {"col_1": pa.string(), "col_2": pa.intaa()}
		assert writer._schema == pa.schema(_lowerCamelCase						,   metadata=writer._schema.metadata							)
		_check_output(output.getvalue()						,   expected_num_chunks=num_examples if writer_batch_size == 1 else 1							)



def  SCREAMING_SNAKE_CASE				(      )		->    Dict:





		with tempfile.TemporaryDirectory() as tmp_dir:
				__lowercase											=     {"col_1": pa.string(), "col_2": pa.intaa()}
				__lowercase											=     os.path.join(_lowerCamelCase						,   'test.arrow'							)
				with ArrowWriter(path=_lowerCamelCase						,   schema=pa.schema(_lowerCamelCase							)							) as writer:
						writer.write_batch({'col_1': ['foo', 'bar'], 'col_2': [1, 2]}							)
						__lowercase											=     writer.finalize()
				assert num_examples == 2
				assert num_bytes > 0
				assert writer._schema == pa.schema(_lowerCamelCase						,   metadata=writer._schema.metadata							)
				_check_output(_lowerCamelCase						,   1							)



def  SCREAMING_SNAKE_CASE				(      snake_case							)		->    List[str]:





		if pa.types.is_list(_lowerCamelCase							):
				return get_base_dtype(arr_type.value_type							)
		else:
				return arr_type



def  SCREAMING_SNAKE_CASE				(      snake_case						,   snake_case							)		->    int:





		if isinstance(lst[0]						,   _lowerCamelCase							):
				change_first_primitive_element_in_list(lst[0]						,   _lowerCamelCase							)
		else:
				__lowercase											=     value



@pytest.mark.parametrize('optimized_int_type, expected_dtype'						,   [(None, pa.intaa()), (Value('int32'							), pa.intaa())]							)
@pytest.mark.parametrize('sequence'						,   [[1, 2, 3], [[1, 2, 3]], [[[1, 2, 3]]]]							)
def  SCREAMING_SNAKE_CASE				(      snake_case						,   snake_case						,   snake_case							)		->    Any:
		__lowercase											=     pa.array(TypedSequence(_lowerCamelCase						,   optimized_int_type=_lowerCamelCase							)							)
		assert get_base_dtype(arr.type							) == expected_dtype



@pytest.mark.parametrize(
    'col, expected_dtype'						,   [
        ('attention_mask', pa.inta()),
        ('special_tokens_mask', pa.inta()),
        ('token_type_ids', pa.inta()),
        ('input_ids', pa.intaa()),
        ('other', pa.intaa()),
    ]						,   )
@pytest.mark.parametrize('sequence'						,   [[1, 2, 3], [[1, 2, 3]], [[[1, 2, 3]]]]							)
def  SCREAMING_SNAKE_CASE				(      snake_case						,   snake_case						,   snake_case							)		->    Dict:
		__lowercase											=     pa.array(OptimizedTypedSequence(_lowerCamelCase						,   col=_lowerCamelCase							)							)
		assert get_base_dtype(arr.type							) == expected_dtype

		# not in range
		if col != "other":
				# avoids errors due to in-place modifications
				__lowercase											=     copy.deepcopy(_lowerCamelCase							)
				__lowercase											=     np.iinfo(expected_dtype.to_pandas_dtype()							).max + 1
				change_first_primitive_element_in_list(_lowerCamelCase						,   _lowerCamelCase							)
				__lowercase											=     pa.array(OptimizedTypedSequence(_lowerCamelCase						,   col=_lowerCamelCase							)							)
				assert get_base_dtype(arr.type							) == pa.intaa()



@pytest.mark.parametrize('raise_exception'						,   [False, True]							)
def  SCREAMING_SNAKE_CASE				(      snake_case						,   snake_case							)		->    Union[str, Any]:
		__lowercase											=     str(tmp_path / 'dataset-train.arrow'							)
		try:
				with ArrowWriter(path=_lowerCamelCase							) as writer:
						if raise_exception:
								raise pa.lib.ArrowInvalid()
						else:
								writer.stream.close()
		except pa.lib.ArrowInvalid:
				pass
		finally:
				assert writer.stream.closed



def  SCREAMING_SNAKE_CASE				(      snake_case							)		->    Any:
		__lowercase											=     "mock://dataset-train.arrow"
		with ArrowWriter(path=_lowerCamelCase						,   storage_options=mockfs.storage_options							) as writer:
				assert isinstance(writer._fs						,   type(_lowerCamelCase							)							)
				assert writer._fs.storage_options == mockfs.storage_options
				writer.write({'col_1': 'foo', 'col_2': 1}							)
				writer.write({'col_1': 'bar', 'col_2': 2}							)
				__lowercase											=     writer.finalize()
		assert num_examples == 2
		assert num_bytes > 0
		assert mockfs.exists(_lowerCamelCase							)



def  SCREAMING_SNAKE_CASE				(      )		->    Optional[int]:
		__lowercase											=     pa.BufferOutputStream()
		with ParquetWriter(stream=_lowerCamelCase							) as writer:
				writer.write({'col_1': 'foo', 'col_2': 1}							)
				writer.write({'col_1': 'bar', 'col_2': 2}							)
				__lowercase											=     writer.finalize()
		assert num_examples == 2
		assert num_bytes > 0
		__lowercase											=     pa.BufferReader(output.getvalue()							)
		__lowercase											=     pq.read_table(_lowerCamelCase							)
		assert pa_table.to_pydict() == {"col_1": ["foo", "bar"], "col_2": [1, 2]}



@require_pil
@pytest.mark.parametrize('embed_local_files'						,   [False, True]							)
def  SCREAMING_SNAKE_CASE				(      snake_case						,   snake_case							)		->    Dict:
		import PIL.Image

		__lowercase											=     str(tmp_path / 'test_image_rgb.jpg'							)
		PIL.Image.fromarray(np.zeros((5, 5)						,   dtype=np.uinta							)							).save(_lowerCamelCase						,   format='png'							)
		__lowercase											=     pa.BufferOutputStream()
		with ParquetWriter(
		    stream=_lowerCamelCase						,   features=Features({'image': Image()}							)						,   embed_local_files=_lowerCamelCase							) as writer:
				writer.write({'image': image_path}							)
				writer.finalize()
		__lowercase											=     pa.BufferReader(output.getvalue()							)
		__lowercase											=     pq.read_table(_lowerCamelCase							)
		__lowercase											=     pa_table.to_pydict()
		if embed_local_files:
				assert isinstance(out['image'][0]['path']						,   _lowerCamelCase							)
				with open(_lowerCamelCase						,   'rb'							) as f:
						assert out["image"][0]["bytes"] == f.read()
		else:
				assert out["image"][0]["path"] == image_path
				assert out["image"][0]["bytes"] is None



def  SCREAMING_SNAKE_CASE				(      )		->    List[Any]:
		__lowercase											=     pa.schema([pa.field('col_1'						,   pa.string()						,   nullable=_lowerCamelCase							)]							)
		__lowercase											=     pa.BufferOutputStream()
		with ArrowWriter(stream=_lowerCamelCase							) as writer:
				writer._build_writer(inferred_schema=_lowerCamelCase							)
				assert writer._schema == pa.schema([pa.field('col_1'						,   pa.string()							)]							)






import unittest

from transformers.testing_utils import require_bsa
from transformers.utils import is_bsa_available

from ...test_feature_extraction_common import FeatureExtractionSavingTestMixin


if is_bsa_available():
		from transformers import MarkupLMFeatureExtractor




class 					lowerCamelCase__				(							unittest.TestCase):




	'''simple docstring'''




	def __init__(   self						:Optional[int]	,  a						:List[str]		)						->					Union[str, Any]:
								__UpperCamelCase      :  Dict       =   parent







	def   _lowerCamelCase							(   self						:Dict		)						->					int:
								return {}




def       _SCREAMING_SNAKE_CASE   (						)   ->       Optional[Any]:






							'''simple docstring'''





							__UpperCamelCase      :  Union[str, Any]       =   "<HTML>\n\n    <HEAD>\n    <TITLE>sample document</TITLE>\n    </HEAD>\n\n    <BODY BGCOLOR=\"FFFFFF\">\n    <HR>\n    <a href=\"http://google.com\">Goog</a>\n    <H1>This is one header</H1>\n    <H2>This is a another Header</H2>\n    <P>Travel from\n        <P>\n        <B>SFO to JFK</B>\n        <BR>\n        <B><I>on May 2, 2015 at 2:00 pm. For details go to confirm.com </I></B>\n        <HR>\n        <div style=\"color:#0000FF\">\n            <h3>Traveler <b> name </b> is\n            <p> John Doe </p>\n        </div>"

							__UpperCamelCase      :  Optional[int]       =   "\n    <!DOCTYPE html>\n    <html>\n    <body>\n\n    <h1>My First Heading</h1>\n    <p>My first paragraph.</p>\n\n    </body>\n    </html>\n    "

							return [html_string_a, html_string_a]




@require_bsa
class 					lowerCamelCase__				(							__lowercase	,		unittest.TestCase):




	'''simple docstring'''




	_A								=	MarkupLMFeatureExtractor if is_bsa_available() else None







	def   _lowerCamelCase							(   self						:List[str]		)						->					str:
								__UpperCamelCase      :  List[str]       =   MarkupLMFeatureExtractionTester(self		)







	@property
	def   _lowerCamelCase							(   self						:List[Any]		)						->					List[str]:
								return self.feature_extract_tester.prepare_feat_extract_dict()







	def   _lowerCamelCase							(   self						:Dict		)						->					Dict:
								# Initialize feature_extractor
								__UpperCamelCase      :  Tuple       =   self.feature_extraction_class()

								# Test not batched input
								__UpperCamelCase      :  Union[str, Any]       =   get_html_strings()[0]
								__UpperCamelCase      :  Any       =   feature_extractor(a		)

								# fmt: off
								__UpperCamelCase      :  Union[str, Any]       =   [["sample document", "Goog", "This is one header", "This is a another Header", "Travel from", "SFO to JFK", "on May 2, 2015 at 2:00 pm. For details go to confirm.com", "Traveler", "name", "is", "John Doe"]]
								__UpperCamelCase      :  Tuple       =   [["/html/head/title", "/html/body/a", "/html/body/h1", "/html/body/h2", "/html/body/p", "/html/body/p/p/b[1]", "/html/body/p/p/b[2]/i", "/html/body/p/p/div/h3", "/html/body/p/p/div/h3/b", "/html/body/p/p/div/h3", "/html/body/p/p/div/h3/p"]]
								# fmt: on

								self.assertEqual(encoding.nodes	,  a		)
								self.assertEqual(encoding.xpaths	,  a		)

								# Test batched
								__UpperCamelCase      :  str       =   get_html_strings()
								__UpperCamelCase      :  List[Any]       =   feature_extractor(a		)

								# fmt: off
								__UpperCamelCase      :  Optional[int]       =   expected_nodes + [["My First Heading", "My first paragraph."]]
								__UpperCamelCase      :  Tuple       =   expected_xpaths + [["/html/body/h1", "/html/body/p"]]

								self.assertEqual(len(encoding.nodes		)	,  2		)
								self.assertEqual(len(encoding.xpaths		)	,  2		)

								self.assertEqual(encoding.nodes	,  a		)
								self.assertEqual(encoding.xpaths	,  a		)