



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    UpperCAmelCase__(				A__						):

   '''simple docstring'''




   def       UpperCAmelCase		(    self	:      str)						->					str:




         """simple docstring"""

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


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




         """simple docstring"""

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


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




         """simple docstring"""

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


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




         """simple docstring"""

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


   def       UpperCAmelCase		(    self	:      str)						->					int:




         """simple docstring"""

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


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




         """simple docstring"""

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


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




         """simple docstring"""

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


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




         """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	:      Optional[int])						->					int:




         """simple docstring"""

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


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




         """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	:      List[str])						->					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)						->					Tuple:




         """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 	_lowerCAmelCase (    A__    ,	A__						):
      lowercase__  =	pa.BufferReader(__lowercase						) if isinstance(__lowercase    ,	pa.Buffer						) else pa.memory_map(__lowercase						)
      lowercase__  =	pa.ipc.open_stream(__lowercase						)
      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 	_lowerCAmelCase (    A__    ,	A__						):
      lowercase__  =	pa.BufferOutputStream()
      lowercase__  =	pa.schema(__lowercase						) if fields else None
      with ArrowWriter(stream=__lowercase    ,	schema=__lowercase    ,	writer_batch_size=__lowercase						) 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(__lowercase    ,	metadata=writer._schema.metadata						)
      _check_output(output.getvalue()    ,	expected_num_chunks=num_examples if writer_batch_size == 1 else 1						)

def 	_lowerCAmelCase (    ):
      lowercase__  =	pa.BufferOutputStream()
      lowercase__  =	Features({'labels': ClassLabel(names=['neg', 'pos']						)}						)
      with ArrowWriter(stream=__lowercase    ,	features=__lowercase						) 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(__lowercase						)
      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(__lowercase						)

@pytest.mark.parametrize('writer_batch_size'    ,	[None, 1, 10]						)
def 	_lowerCAmelCase (    A__						):
      lowercase__  =	pa.BufferOutputStream()
      with ArrowWriter(
          stream=__lowercase    ,	writer_batch_size=__lowercase    ,	hash_salt='split_name'    ,	check_duplicates=__lowercase    ,	) as writer:
            with pytest.raises(__lowercase						):
                  writer.write({'col_1': 'foo', 'col_2': 1}    ,	key=[1, 2]						)
                  lowercase__  =	writer.finalize()

@pytest.mark.parametrize('writer_batch_size'    ,	[None, 2, 10]						)
def 	_lowerCAmelCase (    A__						):
      lowercase__  =	pa.BufferOutputStream()
      with ArrowWriter(
          stream=__lowercase    ,	writer_batch_size=__lowercase    ,	hash_salt='split_name'    ,	check_duplicates=__lowercase    ,	) as writer:
            with pytest.raises(__lowercase						):
                  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 	_lowerCAmelCase (    A__						):
      lowercase__  =	pa.BufferOutputStream()
      with ArrowWriter(
          stream=__lowercase    ,	writer_batch_size=__lowercase    ,	hash_salt='split_name'    ,	check_duplicates=__lowercase    ,	) 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 	_lowerCAmelCase (    A__    ,	A__						):
      lowercase__  =	pa.BufferOutputStream()
      lowercase__  =	pa.schema(__lowercase						) if fields else None
      with ArrowWriter(stream=__lowercase    ,	schema=__lowercase    ,	writer_batch_size=__lowercase						) 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(__lowercase    ,	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 	_lowerCAmelCase (    A__    ,	A__						):
      lowercase__  =	pa.BufferOutputStream()
      lowercase__  =	pa.schema(__lowercase						) if fields else None
      with ArrowWriter(stream=__lowercase    ,	schema=__lowercase    ,	writer_batch_size=__lowercase						) 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(__lowercase    ,	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 	_lowerCAmelCase (    A__    ,	A__						):
      lowercase__  =	pa.BufferOutputStream()
      lowercase__  =	pa.schema(__lowercase						) if fields else None
      with ArrowWriter(stream=__lowercase    ,	schema=__lowercase    ,	writer_batch_size=__lowercase						) 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(__lowercase    ,	metadata=writer._schema.metadata						)
      _check_output(output.getvalue()    ,	expected_num_chunks=num_examples if writer_batch_size == 1 else 1						)

def 	_lowerCAmelCase (    ):
      with tempfile.TemporaryDirectory() as tmp_dir:
            lowercase__  =	{'col_1': pa.string(), 'col_2': pa.intaa()}
            lowercase__  =	os.path.join(__lowercase    ,	'test.arrow'						)
            with ArrowWriter(path=__lowercase    ,	schema=pa.schema(__lowercase						)						) 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(__lowercase    ,	metadata=writer._schema.metadata						)
            _check_output(__lowercase    ,	1						)

def 	_lowerCAmelCase (    A__						):
      if pa.types.is_list(__lowercase						):
            return get_base_dtype(arr_type.value_type						)
      else:
            return arr_type

def 	_lowerCAmelCase (    A__    ,	A__						):
      if isinstance(lst[0]    ,	__lowercase						):
            change_first_primitive_element_in_list(lst[0]    ,	__lowercase						)
      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 	_lowerCAmelCase (    A__    ,	A__    ,	A__						):
      lowercase__  =	pa.array(TypedSequence(__lowercase    ,	optimized_int_type=__lowercase						)						)
      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 	_lowerCAmelCase (    A__    ,	A__    ,	A__						):
      # in range
      lowercase__  =	pa.array(OptimizedTypedSequence(__lowercase    ,	col=__lowercase						)						)
      assert get_base_dtype(arr.type						) == expected_dtype

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

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

def 	_lowerCAmelCase (    A__						):
      lowercase__  =	'mock://dataset-train.arrow'
      with ArrowWriter(path=__lowercase    ,	storage_options=mockfs.storage_options						) as writer:
            assert isinstance(writer._fs    ,	type(__lowercase						)						)
            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(__lowercase						)

def 	_lowerCAmelCase (    ):
      lowercase__  =	pa.BufferOutputStream()
      with ParquetWriter(stream=__lowercase						) 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(__lowercase						)
      assert pa_table.to_pydict() == {"col_1": ["foo", "bar"], "col_2": [1, 2]}

@require_pil
@pytest.mark.parametrize('embed_local_files'    ,	[False, True]						)
def 	_lowerCAmelCase (    A__    ,	A__						):
      import PIL.Image

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


"""simple docstring"""






from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


__SCREAMING_SNAKE_CASE							=     {
    'configuration_whisper': ['WHISPER_PRETRAINED_CONFIG_ARCHIVE_MAP', 'WhisperConfig', 'WhisperOnnxConfig'],
    'feature_extraction_whisper': ['WhisperFeatureExtractor'],
    'processing_whisper': ['WhisperProcessor'],
    'tokenization_whisper': ['WhisperTokenizer'],
}

try:
 if not is_tokenizers_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 __SCREAMING_SNAKE_CASE							=     ['WhisperTokenizerFast']

try:
 if not is_torch_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 __SCREAMING_SNAKE_CASE							=     [
     'WHISPER_PRETRAINED_MODEL_ARCHIVE_LIST',
     'WhisperForConditionalGeneration',
     'WhisperModel',
     'WhisperPreTrainedModel',
     'WhisperForAudioClassification',
 ]

try:
 if not is_tf_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 __SCREAMING_SNAKE_CASE							=     [
     'TF_WHISPER_PRETRAINED_MODEL_ARCHIVE_LIST',
     'TFWhisperForConditionalGeneration',
     'TFWhisperModel',
     'TFWhisperPreTrainedModel',
 ]

try:
 if not is_flax_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 __SCREAMING_SNAKE_CASE							=     [
     'FlaxWhisperForConditionalGeneration',
     'FlaxWhisperModel',
     'FlaxWhisperPreTrainedModel',
     'FlaxWhisperForAudioClassification',
 ]


if TYPE_CHECKING:
 from .configuration_whisper import WHISPER_PRETRAINED_CONFIG_ARCHIVE_MAP, WhisperConfig, WhisperOnnxConfig
 from .feature_extraction_whisper import WhisperFeatureExtractor
 from .processing_whisper import WhisperProcessor
 from .tokenization_whisper import WhisperTokenizer

 try:
  if not is_tokenizers_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .tokenization_whisper_fast import WhisperTokenizerFast

 try:
  if not is_torch_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_whisper import (
      WHISPER_PRETRAINED_MODEL_ARCHIVE_LIST,
      WhisperForAudioClassification,
      WhisperForConditionalGeneration,
      WhisperModel,
      WhisperPreTrainedModel,
  )

 try:
  if not is_tf_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_tf_whisper import (
      TF_WHISPER_PRETRAINED_MODEL_ARCHIVE_LIST,
      TFWhisperForConditionalGeneration,
      TFWhisperModel,
      TFWhisperPreTrainedModel,
  )

 try:
  if not is_flax_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_flax_whisper import (
      FlaxWhisperForAudioClassification,
      FlaxWhisperForConditionalGeneration,
      FlaxWhisperModel,
      FlaxWhisperPreTrainedModel,
  )

else:
 import sys

 __SCREAMING_SNAKE_CASE							=     _LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)
