



"""simple docstring"""


import io
import json

import fsspec
import pytest

from datasets import Dataset, DatasetDict, Features, NamedSplit, Value
from datasets.io.json import JsonDatasetReader, JsonDatasetWriter

from ..utils import assert_arrow_memory_doesnt_increase, assert_arrow_memory_increases







def  lowercase	(      _SCREAMING_SNAKE_CASE   :						str					,     _SCREAMING_SNAKE_CASE   :						List[Any]				):






				'''simple docstring'''



				assert isinstance(_SCREAMING_SNAKE_CASE					,     _SCREAMING_SNAKE_CASE				)
				assert dataset.num_rows == 4
				assert dataset.num_columns == 3
				assert dataset.column_names == ["col_1", "col_2", "col_3"]
				for feature, expected_dtype in expected_features.items():
								assert dataset.features[feature].dtype == expected_dtype







@pytest.mark.parametrize('''keep_in_memory'''					,     [False, True]				)
def  lowercase	(      _SCREAMING_SNAKE_CASE   :						Tuple					,     _SCREAMING_SNAKE_CASE   :						Dict					,     _SCREAMING_SNAKE_CASE   :						Optional[Any]				):






				'''simple docstring'''



				_UpperCAmelCase					=    tmp_path / '''cache'''
				_UpperCAmelCase					=    {'''col_1''': '''string''', '''col_2''': '''int64''', '''col_3''': '''float64'''}
				with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
								_UpperCAmelCase					=    JsonDatasetReader(_SCREAMING_SNAKE_CASE					,     cache_dir=_SCREAMING_SNAKE_CASE					,     keep_in_memory=_SCREAMING_SNAKE_CASE				).read()
				_check_json_dataset(_SCREAMING_SNAKE_CASE					,     _SCREAMING_SNAKE_CASE				)







@pytest.mark.parametrize(
    '''features'''					,     [
        None,
        {'''col_1''': '''string''', '''col_2''': '''int64''', '''col_3''': '''float64'''},
        {'''col_1''': '''string''', '''col_2''': '''string''', '''col_3''': '''string'''},
        {'''col_1''': '''int32''', '''col_2''': '''int32''', '''col_3''': '''int32'''},
        {'''col_1''': '''float32''', '''col_2''': '''float32''', '''col_3''': '''float32'''},
    ]					,     )
def  lowercase	(      _SCREAMING_SNAKE_CASE   :						Optional[int]					,     _SCREAMING_SNAKE_CASE   :						List[str]					,     _SCREAMING_SNAKE_CASE   :						List[str]				):






				'''simple docstring'''



				_UpperCAmelCase					=    tmp_path / '''cache'''
				_UpperCAmelCase					=    {'''col_1''': '''string''', '''col_2''': '''int64''', '''col_3''': '''float64'''}
				_UpperCAmelCase					=    features.copy() if features else default_expected_features
				_UpperCAmelCase					=    (
				    Features({feature: Value(_SCREAMING_SNAKE_CASE				) for feature, dtype in features.items()}				) if features is not None else None
				)
				_UpperCAmelCase					=    JsonDatasetReader(_SCREAMING_SNAKE_CASE					,     features=_SCREAMING_SNAKE_CASE					,     cache_dir=_SCREAMING_SNAKE_CASE				).read()
				_check_json_dataset(_SCREAMING_SNAKE_CASE					,     _SCREAMING_SNAKE_CASE				)







@pytest.mark.parametrize(
    '''features'''					,     [
        None,
        {'''col_3''': '''float64''', '''col_1''': '''string''', '''col_2''': '''int64'''},
    ]					,     )
def  lowercase	(      _SCREAMING_SNAKE_CASE   :						Union[str, Any]					,     _SCREAMING_SNAKE_CASE   :						Optional[int]					,     _SCREAMING_SNAKE_CASE   :						int				):






				'''simple docstring'''



				_UpperCAmelCase					=    tmp_path / '''cache'''
				_UpperCAmelCase					=    {'''col_3''': '''float64''', '''col_1''': '''string''', '''col_2''': '''int64'''}
				_UpperCAmelCase					=    features.copy() if features else default_expected_features
				_UpperCAmelCase					=    (
				    Features({feature: Value(_SCREAMING_SNAKE_CASE				) for feature, dtype in features.items()}				) if features is not None else None
				)
				_UpperCAmelCase					=    JsonDatasetReader(_SCREAMING_SNAKE_CASE					,     features=_SCREAMING_SNAKE_CASE					,     cache_dir=_SCREAMING_SNAKE_CASE				).read()
				assert isinstance(_SCREAMING_SNAKE_CASE					,     _SCREAMING_SNAKE_CASE				)
				assert dataset.num_rows == 2
				assert dataset.num_columns == 3
				assert dataset.column_names == ["col_3", "col_1", "col_2"]
				for feature, expected_dtype in expected_features.items():
								assert dataset.features[feature].dtype == expected_dtype







def  lowercase	(      _SCREAMING_SNAKE_CASE   :						Dict					,     _SCREAMING_SNAKE_CASE   :						Any				):






				'''simple docstring'''



				_UpperCAmelCase					=    {'''col_2''': '''int64''', '''col_3''': '''float64''', '''col_1''': '''string'''}
				_UpperCAmelCase					=    features.copy()
				_UpperCAmelCase					=    (
				    Features({feature: Value(_SCREAMING_SNAKE_CASE				) for feature, dtype in features.items()}				) if features is not None else None
				)
				_UpperCAmelCase					=    tmp_path / '''cache'''
				_UpperCAmelCase					=    JsonDatasetReader(_SCREAMING_SNAKE_CASE					,     features=_SCREAMING_SNAKE_CASE					,     cache_dir=_SCREAMING_SNAKE_CASE				).read()
				assert isinstance(_SCREAMING_SNAKE_CASE					,     _SCREAMING_SNAKE_CASE				)
				assert dataset.num_rows == 2
				assert dataset.num_columns == 3
				assert dataset.column_names == ["col_2", "col_3", "col_1"]
				for feature, expected_dtype in expected_features.items():
								assert dataset.features[feature].dtype == expected_dtype







@pytest.mark.parametrize('''split'''					,     [None, NamedSplit('''train'''				), '''train''', '''test''']				)
def  lowercase	(      _SCREAMING_SNAKE_CASE   :						List[str]					,     _SCREAMING_SNAKE_CASE   :						List[str]					,     _SCREAMING_SNAKE_CASE   :						Tuple				):






				'''simple docstring'''



				_UpperCAmelCase					=    tmp_path / '''cache'''
				_UpperCAmelCase					=    {'''col_1''': '''string''', '''col_2''': '''int64''', '''col_3''': '''float64'''}
				_UpperCAmelCase					=    JsonDatasetReader(_SCREAMING_SNAKE_CASE					,     cache_dir=_SCREAMING_SNAKE_CASE					,     split=_SCREAMING_SNAKE_CASE				).read()
				_check_json_dataset(_SCREAMING_SNAKE_CASE					,     _SCREAMING_SNAKE_CASE				)
				assert dataset.split == split if split else "train"







@pytest.mark.parametrize('''path_type'''					,     [str, list]				)
def  lowercase	(      _SCREAMING_SNAKE_CASE   :						Dict					,     _SCREAMING_SNAKE_CASE   :						List[str]					,     _SCREAMING_SNAKE_CASE   :						Dict				):






				'''simple docstring'''



				if issubclass(_SCREAMING_SNAKE_CASE					,     _SCREAMING_SNAKE_CASE				):
								_UpperCAmelCase					=    jsonl_path
				elif issubclass(_SCREAMING_SNAKE_CASE					,     _SCREAMING_SNAKE_CASE				):
								_UpperCAmelCase					=    [jsonl_path]
				_UpperCAmelCase					=    tmp_path / '''cache'''
				_UpperCAmelCase					=    {'''col_1''': '''string''', '''col_2''': '''int64''', '''col_3''': '''float64'''}
				_UpperCAmelCase					=    JsonDatasetReader(_SCREAMING_SNAKE_CASE					,     cache_dir=_SCREAMING_SNAKE_CASE				).read()
				_check_json_dataset(_SCREAMING_SNAKE_CASE					,     _SCREAMING_SNAKE_CASE				)







def  lowercase	(      _SCREAMING_SNAKE_CASE   :						Dict					,     _SCREAMING_SNAKE_CASE   :						Optional[int]					,     _SCREAMING_SNAKE_CASE   :						List[Any]=("train",)				):






				'''simple docstring'''



				assert isinstance(_SCREAMING_SNAKE_CASE					,     _SCREAMING_SNAKE_CASE				)
				for split in splits:
								_UpperCAmelCase					=    dataset_dict[split]
								assert dataset.num_rows == 4
								assert dataset.num_columns == 3
								assert dataset.column_names == ["col_1", "col_2", "col_3"]
								for feature, expected_dtype in expected_features.items():
												assert dataset.features[feature].dtype == expected_dtype







@pytest.mark.parametrize('''keep_in_memory'''					,     [False, True]				)
def  lowercase	(      _SCREAMING_SNAKE_CASE   :						Optional[int]					,     _SCREAMING_SNAKE_CASE   :						int					,     _SCREAMING_SNAKE_CASE   :						List[Any]				):






				'''simple docstring'''



				_UpperCAmelCase					=    tmp_path / '''cache'''
				_UpperCAmelCase					=    {'''col_1''': '''string''', '''col_2''': '''int64''', '''col_3''': '''float64'''}
				with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
								_UpperCAmelCase					=    JsonDatasetReader({'''train''': jsonl_path}					,     cache_dir=_SCREAMING_SNAKE_CASE					,     keep_in_memory=_SCREAMING_SNAKE_CASE				).read()
				_check_json_datasetdict(_SCREAMING_SNAKE_CASE					,     _SCREAMING_SNAKE_CASE				)







@pytest.mark.parametrize(
    '''features'''					,     [
        None,
        {'''col_1''': '''string''', '''col_2''': '''int64''', '''col_3''': '''float64'''},
        {'''col_1''': '''string''', '''col_2''': '''string''', '''col_3''': '''string'''},
        {'''col_1''': '''int32''', '''col_2''': '''int32''', '''col_3''': '''int32'''},
        {'''col_1''': '''float32''', '''col_2''': '''float32''', '''col_3''': '''float32'''},
    ]					,     )
def  lowercase	(      _SCREAMING_SNAKE_CASE   :						List[str]					,     _SCREAMING_SNAKE_CASE   :						str					,     _SCREAMING_SNAKE_CASE   :						Tuple				):






				'''simple docstring'''



				_UpperCAmelCase					=    tmp_path / '''cache'''
				_UpperCAmelCase					=    {'''col_1''': '''string''', '''col_2''': '''int64''', '''col_3''': '''float64'''}
				_UpperCAmelCase					=    features.copy() if features else default_expected_features
				_UpperCAmelCase					=    (
				    Features({feature: Value(_SCREAMING_SNAKE_CASE				) for feature, dtype in features.items()}				) if features is not None else None
				)
				_UpperCAmelCase					=    JsonDatasetReader({'''train''': jsonl_path}					,     features=_SCREAMING_SNAKE_CASE					,     cache_dir=_SCREAMING_SNAKE_CASE				).read()
				_check_json_datasetdict(_SCREAMING_SNAKE_CASE					,     _SCREAMING_SNAKE_CASE				)







@pytest.mark.parametrize('''split'''					,     [None, NamedSplit('''train'''				), '''train''', '''test''']				)
def  lowercase	(      _SCREAMING_SNAKE_CASE   :						List[str]					,     _SCREAMING_SNAKE_CASE   :						int					,     _SCREAMING_SNAKE_CASE   :						str				):






				'''simple docstring'''



				if split:
								_UpperCAmelCase					=    {split: jsonl_path}
				else:
								_UpperCAmelCase					=    '''train'''
								_UpperCAmelCase					=    {'''train''': jsonl_path, '''test''': jsonl_path}
				_UpperCAmelCase					=    tmp_path / '''cache'''
				_UpperCAmelCase					=    {'''col_1''': '''string''', '''col_2''': '''int64''', '''col_3''': '''float64'''}
				_UpperCAmelCase					=    JsonDatasetReader(_SCREAMING_SNAKE_CASE					,     cache_dir=_SCREAMING_SNAKE_CASE				).read()
				_check_json_datasetdict(_SCREAMING_SNAKE_CASE					,     _SCREAMING_SNAKE_CASE					,     splits=list(path.keys()				)				)
				assert all(dataset[split].split == split for split in path.keys()				)







def  lowercase	(      _SCREAMING_SNAKE_CASE   :						Union[str, Any]				):






				'''simple docstring'''



				return json.load(_SCREAMING_SNAKE_CASE				)







def  lowercase	(      _SCREAMING_SNAKE_CASE   :						Tuple				):






				'''simple docstring'''



				return [json.loads(_SCREAMING_SNAKE_CASE				) for line in buffer]
class 			_a :
			"""simple docstring"""



			@pytest.mark.parametrize('''lines, load_json_function'''				,   [(True, load_json_lines), (False, load_json)]   )
			def  lowercase__						(		self     :    Dict				,   __UpperCamelCase     :    List[str]				,   __UpperCamelCase     :    int				,   __UpperCamelCase     :    Union[str, Any]   )->Optional[Any]:
							with io.BytesIO() as buffer:
											JsonDatasetWriter(__UpperCamelCase				,   __UpperCamelCase				,   lines=__UpperCamelCase   ).write()
											buffer.seek(0   )
											_UpperCAmelCase					=    load_json_function(__UpperCamelCase   )
							assert isinstance(__UpperCamelCase				,   __UpperCamelCase   )
							assert isinstance(exported_content[0]				,   __UpperCamelCase   )
							assert len(__UpperCamelCase   ) == 1_0



			@pytest.mark.parametrize(
			    '''orient, container, keys, len_at'''				,   [
			        ('''records''', list, {'''tokens''', '''labels''', '''answers''', '''id'''}, None),
			        ('''split''', dict, {'''columns''', '''data'''}, '''data'''),
			        ('''index''', dict, set('''0123456789'''   ), None),
			        ('''columns''', dict, {'''tokens''', '''labels''', '''answers''', '''id'''}, '''tokens'''),
			        ('''values''', list, None, None),
			        ('''table''', dict, {'''schema''', '''data'''}, '''data'''),
			    ]				,   )
			def  lowercase__						(		self     :    int				,   __UpperCamelCase     :    Tuple				,   __UpperCamelCase     :    Dict				,   __UpperCamelCase     :    Optional[int]				,   __UpperCamelCase     :    Any				,   __UpperCamelCase     :    List[str]   )->List[str]:
							with io.BytesIO() as buffer:
											JsonDatasetWriter(__UpperCamelCase				,   __UpperCamelCase				,   lines=__UpperCamelCase				,   orient=__UpperCamelCase   ).write()
											buffer.seek(0   )
											_UpperCAmelCase					=    load_json(__UpperCamelCase   )
							assert isinstance(__UpperCamelCase				,   __UpperCamelCase   )
							if keys:
											if container is dict:
															assert exported_content.keys() == keys
											else:
															assert exported_content[0].keys() == keys
							else:
											assert not hasattr(__UpperCamelCase				,   '''keys'''   ) and not hasattr(exported_content[0]				,   '''keys'''   )
							if len_at:
											assert len(exported_content[len_at]   ) == 1_0
							else:
											assert len(__UpperCamelCase   ) == 1_0



			@pytest.mark.parametrize('''lines, load_json_function'''				,   [(True, load_json_lines), (False, load_json)]   )
			def  lowercase__						(		self     :    int				,   __UpperCamelCase     :    List[Any]				,   __UpperCamelCase     :    int				,   __UpperCamelCase     :    int   )->Any:
							with io.BytesIO() as buffer:
											JsonDatasetWriter(__UpperCamelCase				,   __UpperCamelCase				,   lines=__UpperCamelCase				,   num_proc=2   ).write()
											buffer.seek(0   )
											_UpperCAmelCase					=    load_json_function(__UpperCamelCase   )
							assert isinstance(__UpperCamelCase				,   __UpperCamelCase   )
							assert isinstance(exported_content[0]				,   __UpperCamelCase   )
							assert len(__UpperCamelCase   ) == 1_0



			@pytest.mark.parametrize(
			    '''orient, container, keys, len_at'''				,   [
			        ('''records''', list, {'''tokens''', '''labels''', '''answers''', '''id'''}, None),
			        ('''split''', dict, {'''columns''', '''data'''}, '''data'''),
			        ('''index''', dict, set('''0123456789'''   ), None),
			        ('''columns''', dict, {'''tokens''', '''labels''', '''answers''', '''id'''}, '''tokens'''),
			        ('''values''', list, None, None),
			        ('''table''', dict, {'''schema''', '''data'''}, '''data'''),
			    ]				,   )
			def  lowercase__						(		self     :    Optional[int]				,   __UpperCamelCase     :    List[Any]				,   __UpperCamelCase     :    int				,   __UpperCamelCase     :    Dict				,   __UpperCamelCase     :    Optional[Any]				,   __UpperCamelCase     :    Any   )->int:
							with io.BytesIO() as buffer:
											JsonDatasetWriter(__UpperCamelCase				,   __UpperCamelCase				,   lines=__UpperCamelCase				,   orient=__UpperCamelCase				,   num_proc=2   ).write()
											buffer.seek(0   )
											_UpperCAmelCase					=    load_json(__UpperCamelCase   )
							assert isinstance(__UpperCamelCase				,   __UpperCamelCase   )
							if keys:
											if container is dict:
															assert exported_content.keys() == keys
											else:
															assert exported_content[0].keys() == keys
							else:
											assert not hasattr(__UpperCamelCase				,   '''keys'''   ) and not hasattr(exported_content[0]				,   '''keys'''   )
							if len_at:
											assert len(exported_content[len_at]   ) == 1_0
							else:
											assert len(__UpperCamelCase   ) == 1_0



			def  lowercase__						(		self     :    Optional[Any]				,   __UpperCamelCase     :    Any   )->Optional[int]:
							with pytest.raises(__UpperCamelCase   ):
											with io.BytesIO() as buffer:
															JsonDatasetWriter(__UpperCamelCase				,   __UpperCamelCase				,   num_proc=0   )



			@pytest.mark.parametrize('''compression, extension'''				,   [('''gzip''', '''gz'''), ('''bz2''', '''bz2'''), ('''xz''', '''xz''')]   )
			def  lowercase__						(		self     :    Optional[Any]				,   __UpperCamelCase     :    List[str]				,   __UpperCamelCase     :    int				,   __UpperCamelCase     :    Tuple				,   __UpperCamelCase     :    int				,   __UpperCamelCase     :    Union[str, Any]   )->str:
							_UpperCAmelCase					=    tmp_path_factory.mktemp('''data'''   ) / F'test.json.{extension}'
							_UpperCAmelCase					=    str(shared_datadir / F'test_file.json.{extension}'   )
							JsonDatasetWriter(__UpperCamelCase				,   __UpperCamelCase				,   compression=__UpperCamelCase   ).write()

							with fsspec.open(__UpperCamelCase				,   '''rb'''				,   compression='''infer'''   ) as f:
											_UpperCAmelCase					=    f.read()
							with fsspec.open(__UpperCamelCase				,   '''rb'''				,   compression='''infer'''   ) as f:
											_UpperCAmelCase					=    f.read()
							assert exported_content == original_content








"""simple docstring"""


import inspect
import os
import unittest

import torch

import accelerate
from accelerate import debug_launcher
from accelerate.test_utils import (
    execute_subprocess_async,
    require_cpu,
    require_huggingface_suite,
    require_multi_gpu,
    require_single_gpu,
)
from accelerate.utils import patch_environment
@require_huggingface_suite
class 			_a (					unittest.TestCase):
			"""simple docstring"""



			def  lowercase__						(		self     :    Tuple   )->int:
							_UpperCAmelCase					=    inspect.getfile(accelerate.test_utils   )
							_UpperCAmelCase					=    os.path.sep.join(
							    mod_file.split(os.path.sep   )[:-1] + ['''scripts''', '''external_deps''', '''test_metrics.py''']   )

							from accelerate.test_utils.scripts.external_deps import test_metrics  # noqa: F401

							_UpperCAmelCase					=    test_metrics



			@require_cpu
			def  lowercase__						(		self     :    Any   )->int:
							debug_launcher(self.test_metrics.main				,   num_processes=1   )



			@require_cpu
			def  lowercase__						(		self     :    List[str]   )->List[str]:
							debug_launcher(self.test_metrics.main   )



			@require_single_gpu
			def  lowercase__						(		self     :    List[Any]   )->Dict:
							self.test_metrics.main()



			@require_multi_gpu
			def  lowercase__						(		self     :    str   )->int:
							print(F'Found {torch.cuda.device_count()} devices.'   )
							_UpperCAmelCase					=    ['''torchrun''', F'--nproc_per_node={torch.cuda.device_count()}', self.test_file_path]
							with patch_environment(omp_num_threads=1   ):
											execute_subprocess_async(__UpperCamelCase				,   env=os.environ.copy()   )



