





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_    (		__UpperCamelCase						):






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






			'''simple docstring'''

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






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






			'''simple docstring'''

			with self.assertRaises(_lowerCAmelCase      ):
				__lowerCAmelCase					:	Dict      =       pa.array(TypedSequence([1, 2, 3]      )       ,  type=pa.intaa()      )






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






			'''simple docstring'''

			with self.assertRaises(_lowerCAmelCase      ):
				__lowerCAmelCase					:	str      =       pa.array(TypedSequence([1, 2, 3]       ,  try_type=Value("""bool"""      )       ,  type=Value("""int64"""      )      )      )






		def 				UpperCAmelCase__						(							self			:					int      )->str:






			'''simple docstring'''

			__lowerCAmelCase					:	Optional[int]      =       pa.array(TypedSequence([1, 2, 3]       ,  type=Value("""int32"""      )      )      )
			self.assertEqual(arr.type       ,  pa.intaa()      )






		def 				UpperCAmelCase__						(							self			:					Optional[int]      )->int:






			'''simple docstring'''

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






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






			'''simple docstring'''

			__lowerCAmelCase					:	Any      =       pa.array(TypedSequence([1, 2, 3]       ,  try_type=Value("""int32"""      )      )      )
			self.assertEqual(arr.type       ,  pa.intaa()      )






		def 				UpperCAmelCase__						(							self			:					Optional[Any]      )->str:






			'''simple docstring'''

			__lowerCAmelCase					:	Optional[int]      =       pa.array(TypedSequence(["""foo""", """bar"""]       ,  try_type=Value("""int64"""      )      )      )
			self.assertEqual(arr.type       ,  pa.string()      )






		def 				UpperCAmelCase__						(							self			:					Any      )->List[Any]:






			'''simple docstring'''

			__lowerCAmelCase					:	Optional[Any]      =       pa.array(TypedSequence([[[1, 2, 3]]]       ,  type=ArrayaD((1, 3)       ,  """int64"""      )      )      )
			self.assertEqual(arr.type       ,  ArrayaDExtensionType((1, 3)       ,  """int64"""      )      )






		def 				UpperCAmelCase__						(							self			:					List[Any]      )->Tuple:






			'''simple docstring'''

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






		def 				UpperCAmelCase__						(							self			:					List[Any]      )->List[Any]:






			'''simple docstring'''

			__lowerCAmelCase					:	str      =       pa.array(TypedSequence([[[1, 2, 3]]]       ,  try_type=ArrayaD((1, 3)       ,  """int64"""      )      )      )
			self.assertEqual(arr.type       ,  ArrayaDExtensionType((1, 3)       ,  """int64"""      )      )






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






			'''simple docstring'''

			__lowerCAmelCase					:	List[Any]      =       pa.array(TypedSequence(["""foo""", """bar"""]       ,  try_type=ArrayaD((1, 3)       ,  """int64"""      )      )      )
			self.assertEqual(arr.type       ,  pa.string()      )






		@require_pil
		def 				UpperCAmelCase__						(							self			:					List[str]      )->int:






			'''simple docstring'''

			import PIL.Image

			__lowerCAmelCase					:	Dict      =       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:
				__lowerCAmelCase					:	int      =       pa.array(TypedSequence([{"""path""": None, """bytes""": b"""image_bytes"""}, pil_image]       ,  type=Image()      )      )
				__lowerCAmelCase					:	List[str]      =       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		(      SCREAMING_SNAKE_CASE		:Tuple   ,   SCREAMING_SNAKE_CASE		:int   )      ->   str:
	__lowerCAmelCase					:	Dict      =       pa.BufferReader(__SCREAMING_SNAKE_CASE   ) if isinstance(__SCREAMING_SNAKE_CASE   ,   pa.Buffer   ) else pa.memory_map(__SCREAMING_SNAKE_CASE   )
	__lowerCAmelCase					:	Tuple      =       pa.ipc.open_stream(__SCREAMING_SNAKE_CASE   )
	__lowerCAmelCase					:	pa.Table      =       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		(      SCREAMING_SNAKE_CASE		:str   ,   SCREAMING_SNAKE_CASE		:List[Any]   )      ->   str:
	__lowerCAmelCase					:	Dict      =       pa.BufferOutputStream()
	__lowerCAmelCase					:	List[str]      =       pa.schema(__SCREAMING_SNAKE_CASE   ) if fields else None
	with ArrowWriter(stream=__SCREAMING_SNAKE_CASE   ,   schema=__SCREAMING_SNAKE_CASE   ,   writer_batch_size=__SCREAMING_SNAKE_CASE   ) as writer:
		writer.write({"""col_1""": """foo""", """col_2""": 1}   )
		writer.write({"""col_1""": """bar""", """col_2""": 2}   )
		__lowerCAmelCase					:	List[str]      =       writer.finalize()
	assert num_examples == 2
	assert num_bytes > 0
	if not fields:
		__lowerCAmelCase					:	List[str]      =       {"""col_1""": pa.string(), """col_2""": pa.intaa()}
	assert writer._schema == pa.schema(__SCREAMING_SNAKE_CASE   ,   metadata=writer._schema.metadata   )
	_check_output(output.getvalue()   ,   expected_num_chunks=num_examples if writer_batch_size == 1 else 1   )






def 	_SCREAMING_SNAKE_CASE		(      )      ->   Optional[int]:
	__lowerCAmelCase					:	Optional[Any]      =       pa.BufferOutputStream()
	__lowerCAmelCase					:	Any      =       Features({"""labels""": ClassLabel(names=["""neg""", """pos"""]   )}   )
	with ArrowWriter(stream=__SCREAMING_SNAKE_CASE   ,   features=__SCREAMING_SNAKE_CASE   ) as writer:
		writer.write({"""labels""": 0}   )
		writer.write({"""labels""": 1}   )
		__lowerCAmelCase					:	Tuple      =       writer.finalize()
	assert num_examples == 2
	assert num_bytes > 0
	assert writer._schema == features.arrow_schema
	assert writer._schema.metadata == features.arrow_schema.metadata
	__lowerCAmelCase					:	int      =       pa.BufferReader(output.getvalue()   )
	__lowerCAmelCase					:	List[Any]      =       pa.ipc.open_stream(__SCREAMING_SNAKE_CASE   )
	__lowerCAmelCase					:	pa.Table      =       f.read_all()
	__lowerCAmelCase					:	int      =       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(__SCREAMING_SNAKE_CASE   )






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






@pytest.mark.parametrize("""writer_batch_size"""   ,   [None, 2, 10]   )
def 	_SCREAMING_SNAKE_CASE		(      SCREAMING_SNAKE_CASE		:Dict   )      ->   Dict:
	__lowerCAmelCase					:	str      =       pa.BufferOutputStream()
	with ArrowWriter(
	    stream=__SCREAMING_SNAKE_CASE   ,   writer_batch_size=__SCREAMING_SNAKE_CASE   ,   hash_salt="""split_name"""   ,   check_duplicates=__SCREAMING_SNAKE_CASE   ,   ) as writer:
		with pytest.raises(__SCREAMING_SNAKE_CASE   ):
			writer.write({"""col_1""": """foo""", """col_2""": 1}   ,   key=10   )
			writer.write({"""col_1""": """bar""", """col_2""": 2}   ,   key=10   )
			__lowerCAmelCase					:	str      =       writer.finalize()






@pytest.mark.parametrize("""writer_batch_size"""   ,   [None, 2, 10]   )
def 	_SCREAMING_SNAKE_CASE		(      SCREAMING_SNAKE_CASE		:Optional[int]   )      ->   Tuple:
	__lowerCAmelCase					:	str      =       pa.BufferOutputStream()
	with ArrowWriter(
	    stream=__SCREAMING_SNAKE_CASE   ,   writer_batch_size=__SCREAMING_SNAKE_CASE   ,   hash_salt="""split_name"""   ,   check_duplicates=__SCREAMING_SNAKE_CASE   ,   ) as writer:
		writer.write({"""col_1""": """foo""", """col_2""": 1}   ,   key=1   )
		writer.write({"""col_1""": """bar""", """col_2""": 2}   ,   key=2   )
		__lowerCAmelCase					:	Optional[Any]      =       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		(      SCREAMING_SNAKE_CASE		:Optional[int]   ,   SCREAMING_SNAKE_CASE		:Any   )      ->   str:
	__lowerCAmelCase					:	Dict      =       pa.BufferOutputStream()
	__lowerCAmelCase					:	Optional[int]      =       pa.schema(__SCREAMING_SNAKE_CASE   ) if fields else None
	with ArrowWriter(stream=__SCREAMING_SNAKE_CASE   ,   schema=__SCREAMING_SNAKE_CASE   ,   writer_batch_size=__SCREAMING_SNAKE_CASE   ) as writer:
		writer.write_batch({"""col_1""": ["""foo""", """bar"""], """col_2""": [1, 2]}   )
		writer.write_batch({"""col_1""": [], """col_2""": []}   )
		__lowerCAmelCase					:	int      =       writer.finalize()
	assert num_examples == 2
	assert num_bytes > 0
	if not fields:
		__lowerCAmelCase					:	List[Any]      =       {"""col_1""": pa.string(), """col_2""": pa.intaa()}
	assert writer._schema == pa.schema(__SCREAMING_SNAKE_CASE   ,   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		(      SCREAMING_SNAKE_CASE		:Union[str, Any]   ,   SCREAMING_SNAKE_CASE		:Any   )      ->   Union[str, Any]:
	__lowerCAmelCase					:	int      =       pa.BufferOutputStream()
	__lowerCAmelCase					:	Any      =       pa.schema(__SCREAMING_SNAKE_CASE   ) if fields else None
	with ArrowWriter(stream=__SCREAMING_SNAKE_CASE   ,   schema=__SCREAMING_SNAKE_CASE   ,   writer_batch_size=__SCREAMING_SNAKE_CASE   ) as writer:
		writer.write_table(pa.Table.from_pydict({"""col_1""": ["""foo""", """bar"""], """col_2""": [1, 2]}   )   )
		__lowerCAmelCase					:	Union[str, Any]      =       writer.finalize()
	assert num_examples == 2
	assert num_bytes > 0
	if not fields:
		__lowerCAmelCase					:	List[Any]      =       {"""col_1""": pa.string(), """col_2""": pa.intaa()}
	assert writer._schema == pa.schema(__SCREAMING_SNAKE_CASE   ,   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		(      SCREAMING_SNAKE_CASE		:List[Any]   ,   SCREAMING_SNAKE_CASE		:str   )      ->   str:
	__lowerCAmelCase					:	str      =       pa.BufferOutputStream()
	__lowerCAmelCase					:	int      =       pa.schema(__SCREAMING_SNAKE_CASE   ) if fields else None
	with ArrowWriter(stream=__SCREAMING_SNAKE_CASE   ,   schema=__SCREAMING_SNAKE_CASE   ,   writer_batch_size=__SCREAMING_SNAKE_CASE   ) 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]}   )   )
		__lowerCAmelCase					:	int      =       writer.finalize()
	assert num_examples == 2
	assert num_bytes > 0
	if not fields:
		__lowerCAmelCase					:	Any      =       {"""col_1""": pa.string(), """col_2""": pa.intaa()}
	assert writer._schema == pa.schema(__SCREAMING_SNAKE_CASE   ,   metadata=writer._schema.metadata   )
	_check_output(output.getvalue()   ,   expected_num_chunks=num_examples if writer_batch_size == 1 else 1   )






def 	_SCREAMING_SNAKE_CASE		(      )      ->   int:
	with tempfile.TemporaryDirectory() as tmp_dir:
		__lowerCAmelCase					:	Union[str, Any]      =       {"""col_1""": pa.string(), """col_2""": pa.intaa()}
		__lowerCAmelCase					:	str      =       os.path.join(__SCREAMING_SNAKE_CASE   ,   """test.arrow"""   )
		with ArrowWriter(path=__SCREAMING_SNAKE_CASE   ,   schema=pa.schema(__SCREAMING_SNAKE_CASE   )   ) as writer:
			writer.write_batch({"""col_1""": ["""foo""", """bar"""], """col_2""": [1, 2]}   )
			__lowerCAmelCase					:	Any      =       writer.finalize()
		assert num_examples == 2
		assert num_bytes > 0
		assert writer._schema == pa.schema(__SCREAMING_SNAKE_CASE   ,   metadata=writer._schema.metadata   )
		_check_output(__SCREAMING_SNAKE_CASE   ,   1   )






def 	_SCREAMING_SNAKE_CASE		(      SCREAMING_SNAKE_CASE		:List[str]   )      ->   Optional[Any]:
	if pa.types.is_list(__SCREAMING_SNAKE_CASE   ):
		return get_base_dtype(arr_type.value_type   )
	else:
		return arr_type






def 	_SCREAMING_SNAKE_CASE		(      SCREAMING_SNAKE_CASE		:List[str]   ,   SCREAMING_SNAKE_CASE		:int   )      ->   Dict:
	if isinstance(lst[0]   ,   __SCREAMING_SNAKE_CASE   ):
		change_first_primitive_element_in_list(lst[0]   ,   __SCREAMING_SNAKE_CASE   )
	else:
		__lowerCAmelCase					:	Tuple      =       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		(      SCREAMING_SNAKE_CASE		:Union[str, Any]   ,   SCREAMING_SNAKE_CASE		:str   ,   SCREAMING_SNAKE_CASE		:Any   )      ->   Tuple:
	__lowerCAmelCase					:	Union[str, Any]      =       pa.array(TypedSequence(__SCREAMING_SNAKE_CASE   ,   optimized_int_type=__SCREAMING_SNAKE_CASE   )   )
	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		(      SCREAMING_SNAKE_CASE		:str   ,   SCREAMING_SNAKE_CASE		:Tuple   ,   SCREAMING_SNAKE_CASE		:Optional[Any]   )      ->   str:
	__lowerCAmelCase					:	Optional[int]      =       pa.array(OptimizedTypedSequence(__SCREAMING_SNAKE_CASE   ,   col=__SCREAMING_SNAKE_CASE   )   )
	assert get_base_dtype(arr.type   ) == expected_dtype

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






@pytest.mark.parametrize("""raise_exception"""   ,   [False, True]   )
def 	_SCREAMING_SNAKE_CASE		(      SCREAMING_SNAKE_CASE		:Dict   ,   SCREAMING_SNAKE_CASE		:Any   )      ->   Optional[Any]:
	__lowerCAmelCase					:	int      =       str(tmp_path / """dataset-train.arrow"""   )
	try:
		with ArrowWriter(path=__SCREAMING_SNAKE_CASE   ) 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		(      SCREAMING_SNAKE_CASE		:Dict   )      ->   Any:
	__lowerCAmelCase					:	Any      =       """mock://dataset-train.arrow"""
	with ArrowWriter(path=__SCREAMING_SNAKE_CASE   ,   storage_options=mockfs.storage_options   ) as writer:
		assert isinstance(writer._fs   ,   type(__SCREAMING_SNAKE_CASE   )   )
		assert writer._fs.storage_options == mockfs.storage_options
		writer.write({"""col_1""": """foo""", """col_2""": 1}   )
		writer.write({"""col_1""": """bar""", """col_2""": 2}   )
		__lowerCAmelCase					:	Tuple      =       writer.finalize()
	assert num_examples == 2
	assert num_bytes > 0
	assert mockfs.exists(__SCREAMING_SNAKE_CASE   )






def 	_SCREAMING_SNAKE_CASE		(      )      ->   Optional[int]:
	__lowerCAmelCase					:	Tuple      =       pa.BufferOutputStream()
	with ParquetWriter(stream=__SCREAMING_SNAKE_CASE   ) as writer:
		writer.write({"""col_1""": """foo""", """col_2""": 1}   )
		writer.write({"""col_1""": """bar""", """col_2""": 2}   )
		__lowerCAmelCase					:	Optional[Any]      =       writer.finalize()
	assert num_examples == 2
	assert num_bytes > 0
	__lowerCAmelCase					:	str      =       pa.BufferReader(output.getvalue()   )
	__lowerCAmelCase					:	pa.Table      =       pq.read_table(__SCREAMING_SNAKE_CASE   )
	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		(      SCREAMING_SNAKE_CASE		:Any   ,   SCREAMING_SNAKE_CASE		:Tuple   )      ->   Optional[int]:
	import PIL.Image

	__lowerCAmelCase					:	List[Any]      =       str(tmp_path / """test_image_rgb.jpg"""   )
	PIL.Image.fromarray(np.zeros((5, 5)   ,   dtype=np.uinta   )   ).save(__SCREAMING_SNAKE_CASE   ,   format="""png"""   )
	__lowerCAmelCase					:	List[str]      =       pa.BufferOutputStream()
	with ParquetWriter(
	    stream=__SCREAMING_SNAKE_CASE   ,   features=Features({"""image""": Image()}   )   ,   embed_local_files=__SCREAMING_SNAKE_CASE   ) as writer:
		writer.write({"""image""": image_path}   )
		writer.finalize()
	__lowerCAmelCase					:	List[str]      =       pa.BufferReader(output.getvalue()   )
	__lowerCAmelCase					:	pa.Table      =       pq.read_table(__SCREAMING_SNAKE_CASE   )
	__lowerCAmelCase					:	List[str]      =       pa_table.to_pydict()
	if embed_local_files:
		assert isinstance(out["""image"""][0]["""path"""]   ,   __SCREAMING_SNAKE_CASE   )
		with open(__SCREAMING_SNAKE_CASE   ,   """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		(      )      ->   Optional[int]:
	__lowerCAmelCase					:	Any      =       pa.schema([pa.field("""col_1"""   ,   pa.string()   ,   nullable=__SCREAMING_SNAKE_CASE   )]   )
	__lowerCAmelCase					:	Union[str, Any]      =       pa.BufferOutputStream()
	with ArrowWriter(stream=__SCREAMING_SNAKE_CASE   ) as writer:
		writer._build_writer(inferred_schema=__SCREAMING_SNAKE_CASE   )
		assert writer._schema == pa.schema([pa.field("""col_1"""   ,   pa.string()   )]   )






import inspect
import unittest

import numpy as np

from transformers import ViTConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_flax_common import FlaxModelTesterMixin, floats_tensor


if is_flax_available():
      import jax

      from transformers.models.vit.modeling_flax_vit import FlaxViTForImageClassification, FlaxViTModel






class        snake_case_    (		unittest.TestCase						):






  def __init__(							self			:					List[Any]       ,  _snake_case			:					List[Any]       ,  _snake_case			:					str=13       ,  _snake_case			:					int=30       ,  _snake_case			:					str=2       ,  _snake_case			:					int=3       ,  _snake_case			:					Optional[Any]=True       ,  _snake_case			:					str=True       ,  _snake_case			:					Optional[int]=32       ,  _snake_case			:					Dict=5       ,  _snake_case			:					Optional[int]=4       ,  _snake_case			:					List[Any]=37       ,  _snake_case			:					Union[str, Any]="gelu"       ,  _snake_case			:					str=0.1       ,  _snake_case			:					str=0.1       ,  _snake_case			:					str=10       ,  _snake_case			:					Any=0.02       ,  )->Tuple:






   '''simple docstring'''

   __lowerCAmelCase					:	Optional[int]      =       parent
   __lowerCAmelCase					:	Any      =       batch_size
   __lowerCAmelCase					:	int      =       image_size
   __lowerCAmelCase					:	int      =       patch_size
   __lowerCAmelCase					:	List[Any]      =       num_channels
   __lowerCAmelCase					:	str      =       is_training
   __lowerCAmelCase					:	str      =       use_labels
   __lowerCAmelCase					:	List[str]      =       hidden_size
   __lowerCAmelCase					:	Dict      =       num_hidden_layers
   __lowerCAmelCase					:	List[str]      =       num_attention_heads
   __lowerCAmelCase					:	Any      =       intermediate_size
   __lowerCAmelCase					:	List[str]      =       hidden_act
   __lowerCAmelCase					:	List[str]      =       hidden_dropout_prob
   __lowerCAmelCase					:	Any      =       attention_probs_dropout_prob
   __lowerCAmelCase					:	Tuple      =       type_sequence_label_size
   __lowerCAmelCase					:	Union[str, Any]      =       initializer_range

   # in ViT, the seq length equals the number of patches + 1 (we add 1 for the [CLS] token)
   __lowerCAmelCase					:	Tuple      =       (image_size // patch_size) ** 2
   __lowerCAmelCase					:	Optional[Any]      =       num_patches + 1






  def 				UpperCAmelCase__						(							self			:					Any      )->Any:






   '''simple docstring'''

   __lowerCAmelCase					:	List[Any]      =       floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]      )

   __lowerCAmelCase					:	Tuple      =       ViTConfig(
       image_size=self.image_size       ,  patch_size=self.patch_size       ,  num_channels=self.num_channels       ,  hidden_size=self.hidden_size       ,  num_hidden_layers=self.num_hidden_layers       ,  num_attention_heads=self.num_attention_heads       ,  intermediate_size=self.intermediate_size       ,  hidden_act=self.hidden_act       ,  hidden_dropout_prob=self.hidden_dropout_prob       ,  attention_probs_dropout_prob=self.attention_probs_dropout_prob       ,  is_decoder=_snake_case       ,  initializer_range=self.initializer_range       ,  )

   return config, pixel_values






  def 				UpperCAmelCase__						(							self			:					List[Any]       ,  _snake_case			:					List[str]       ,  _snake_case			:					List[Any]      )->Any:






   '''simple docstring'''

   __lowerCAmelCase					:	str      =       FlaxViTModel(config=_snake_case      )
   __lowerCAmelCase					:	int      =       model(_snake_case      )
   # expected sequence length = num_patches + 1 (we add 1 for the [CLS] token)
   __lowerCAmelCase					:	Dict      =       (self.image_size, self.image_size)
   __lowerCAmelCase					:	Any      =       (self.patch_size, self.patch_size)
   __lowerCAmelCase					:	Any      =       (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])
   self.parent.assertEqual(result.last_hidden_state.shape       ,  (self.batch_size, num_patches + 1, self.hidden_size)      )






  def 				UpperCAmelCase__						(							self			:					Dict       ,  _snake_case			:					str       ,  _snake_case			:					List[Any]      )->Union[str, Any]:






   '''simple docstring'''

   __lowerCAmelCase					:	Tuple      =       self.type_sequence_label_size
   __lowerCAmelCase					:	Tuple      =       FlaxViTForImageClassification(config=_snake_case      )
   __lowerCAmelCase					:	List[str]      =       model(_snake_case      )
   self.parent.assertEqual(result.logits.shape       ,  (self.batch_size, self.type_sequence_label_size)      )

   # test greyscale images
   __lowerCAmelCase					:	str      =       1
   __lowerCAmelCase					:	Any      =       FlaxViTForImageClassification(_snake_case      )

   __lowerCAmelCase					:	int      =       floats_tensor([self.batch_size, 1, self.image_size, self.image_size]      )
   __lowerCAmelCase					:	Dict      =       model(_snake_case      )






  def 				UpperCAmelCase__						(							self			:					str      )->Any:






   '''simple docstring'''

   __lowerCAmelCase					:	Any      =       self.prepare_config_and_inputs()
   (
       (
       __lowerCAmelCase
   ) ,					(
       __lowerCAmelCase
   ) ,					
   )					:	Union[str, Any]      =       config_and_inputs
   __lowerCAmelCase					:	str      =       {"""pixel_values""": pixel_values}
   return config, inputs_dict






@require_flax
class        snake_case_    (		__lowercase					,unittest.TestCase						):
  A_      							=   (FlaxViTModel, FlaxViTForImageClassification) if is_flax_available() else ()






  def 				UpperCAmelCase__						(							self			:					str      )->None:






   '''simple docstring'''

   __lowerCAmelCase					:	List[str]      =       FlaxViTModelTester(self      )
   __lowerCAmelCase					:	Optional[Any]      =       ConfigTester(self       ,  config_class=_snake_case       ,  has_text_modality=_snake_case       ,  hidden_size=37      )






  def 				UpperCAmelCase__						(							self			:					Any      )->Dict:






   '''simple docstring'''

   self.config_tester.run_common_tests()






  def 				UpperCAmelCase__						(							self			:					Any      )->Any:






   '''simple docstring'''

   __lowerCAmelCase					:	Dict      =       self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_model(*_snake_case      )






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






   '''simple docstring'''

   __lowerCAmelCase					:	Dict      =       self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_for_image_classification(*_snake_case      )






  def 				UpperCAmelCase__						(							self			:					int      )->int:






   '''simple docstring'''

   __lowerCAmelCase ,					__lowerCAmelCase					:	Any      =       self.model_tester.prepare_config_and_inputs_for_common()

   for model_class in self.all_model_classes:
    __lowerCAmelCase					:	int      =       model_class(_snake_case      )
    __lowerCAmelCase					:	Optional[int]      =       inspect.signature(model.__call__      )
    # signature.parameters is an OrderedDict => so arg_names order is deterministic
    __lowerCAmelCase					:	Tuple      =       [*signature.parameters.keys()]

    __lowerCAmelCase					:	str      =       ["""pixel_values"""]
    self.assertListEqual(arg_names[:1]       ,  _snake_case      )






  def 				UpperCAmelCase__						(							self			:					str      )->str:






   '''simple docstring'''

   __lowerCAmelCase ,					__lowerCAmelCase					:	Any      =       self.model_tester.prepare_config_and_inputs_for_common()

   for model_class in self.all_model_classes:
    with self.subTest(model_class.__name__      ):
     __lowerCAmelCase					:	Optional[int]      =       self._prepare_for_class(_snake_case       ,  _snake_case      )
     __lowerCAmelCase					:	List[str]      =       model_class(_snake_case      )

     @jax.jit
     def model_jitted(_snake_case			:					Dict       ,  **_snake_case			:					Union[str, Any]      ):
      return model(pixel_values=_snake_case       ,  **_snake_case      )

     with self.subTest("""JIT Enabled"""      ):
      __lowerCAmelCase					:	List[Any]      =       model_jitted(**_snake_case      ).to_tuple()

     with self.subTest("""JIT Disabled"""      ):
      with jax.disable_jit():
       __lowerCAmelCase					:	str      =       model_jitted(**_snake_case      ).to_tuple()

     self.assertEqual(len(_snake_case      )       ,  len(_snake_case      )      )
     for jitted_output, output in zip(_snake_case       ,  _snake_case      ):
      self.assertEqual(jitted_output.shape       ,  output.shape      )






  @slow
  def 				UpperCAmelCase__						(							self			:					Dict      )->str:






   '''simple docstring'''

   for model_class_name in self.all_model_classes:
    __lowerCAmelCase					:	List[str]      =       model_class_name.from_pretrained("""google/vit-base-patch16-224"""      )
    __lowerCAmelCase					:	List[str]      =       model(np.ones((1, 3, 224, 224)      )      )
    self.assertIsNotNone(_snake_case      )