





'''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 				_lowerCAmelCase  (			lowerCamelCase_     :					Union[str, Any]						,    lowerCamelCase_     :					Union[str, Any]			):
			assert isinstance(lowerCamelCase_						,    lowerCamelCase_			)
			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 				_lowerCAmelCase  (			lowerCamelCase_     :					Union[str, Any]						,    lowerCamelCase_     :					Optional[int]						,    lowerCamelCase_     :					str			):
			__lowercase			  =							tmp_path / '''cache'''
			__lowercase			  =							{'''col_1''': '''string''', '''col_2''': '''int64''', '''col_3''': '''float64'''}
			with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
						__lowercase			  =							JsonDatasetReader(lowerCamelCase_						,    cache_dir=lowerCamelCase_						,    keep_in_memory=lowerCamelCase_			).read()
			_check_json_dataset(lowerCamelCase_						,    lowerCamelCase_			)




@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 				_lowerCAmelCase  (			lowerCamelCase_     :					Optional[Any]						,    lowerCamelCase_     :					Union[str, Any]						,    lowerCamelCase_     :					List[str]			):
			__lowercase			  =							tmp_path / '''cache'''
			__lowercase			  =							{'''col_1''': '''string''', '''col_2''': '''int64''', '''col_3''': '''float64'''}
			__lowercase			  =							features.copy() if features else default_expected_features
			__lowercase			  =							(
			    Features({feature: Value(lowerCamelCase_			) for feature, dtype in features.items()}			) if features is not None else None
			)
			__lowercase			  =							JsonDatasetReader(lowerCamelCase_						,    features=lowerCamelCase_						,    cache_dir=lowerCamelCase_			).read()
			_check_json_dataset(lowerCamelCase_						,    lowerCamelCase_			)




@pytest.mark.parametrize(
    '''features'''						,    [
        None,
        {'''col_3''': '''float64''', '''col_1''': '''string''', '''col_2''': '''int64'''},
    ]						,    )
def 				_lowerCAmelCase  (			lowerCamelCase_     :					Optional[int]						,    lowerCamelCase_     :					Union[str, Any]						,    lowerCamelCase_     :					Union[str, Any]			):
			__lowercase			  =							tmp_path / '''cache'''
			__lowercase			  =							{'''col_3''': '''float64''', '''col_1''': '''string''', '''col_2''': '''int64'''}
			__lowercase			  =							features.copy() if features else default_expected_features
			__lowercase			  =							(
			    Features({feature: Value(lowerCamelCase_			) for feature, dtype in features.items()}			) if features is not None else None
			)
			__lowercase			  =							JsonDatasetReader(lowerCamelCase_						,    features=lowerCamelCase_						,    cache_dir=lowerCamelCase_			).read()
			assert isinstance(lowerCamelCase_						,    lowerCamelCase_			)
			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 				_lowerCAmelCase  (			lowerCamelCase_     :					List[Any]						,    lowerCamelCase_     :					int			):
			# jsonl_312_path features are {"col_3": "float64", "col_1": "string", "col_2": "int64"}
			__lowercase			  =							{'''col_2''': '''int64''', '''col_3''': '''float64''', '''col_1''': '''string'''}
			__lowercase			  =							features.copy()
			__lowercase			  =							(
			    Features({feature: Value(lowerCamelCase_			) for feature, dtype in features.items()}			) if features is not None else None
			)
			__lowercase			  =							tmp_path / '''cache'''
			__lowercase			  =							JsonDatasetReader(lowerCamelCase_						,    features=lowerCamelCase_						,    cache_dir=lowerCamelCase_			).read()
			assert isinstance(lowerCamelCase_						,    lowerCamelCase_			)
			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 				_lowerCAmelCase  (			lowerCamelCase_     :					Optional[int]						,    lowerCamelCase_     :					List[Any]						,    lowerCamelCase_     :					Any			):
			__lowercase			  =							tmp_path / '''cache'''
			__lowercase			  =							{'''col_1''': '''string''', '''col_2''': '''int64''', '''col_3''': '''float64'''}
			__lowercase			  =							JsonDatasetReader(lowerCamelCase_						,    cache_dir=lowerCamelCase_						,    split=lowerCamelCase_			).read()
			_check_json_dataset(lowerCamelCase_						,    lowerCamelCase_			)
			assert dataset.split == split if split else "train"




@pytest.mark.parametrize('''path_type'''						,    [str, list]			)
def 				_lowerCAmelCase  (			lowerCamelCase_     :					Optional[int]						,    lowerCamelCase_     :					Union[str, Any]						,    lowerCamelCase_     :					Tuple			):
			if issubclass(lowerCamelCase_						,    lowerCamelCase_			):
						__lowercase			  =							jsonl_path
			elif issubclass(lowerCamelCase_						,    lowerCamelCase_			):
						__lowercase			  =							[jsonl_path]
			__lowercase			  =							tmp_path / '''cache'''
			__lowercase			  =							{'''col_1''': '''string''', '''col_2''': '''int64''', '''col_3''': '''float64'''}
			__lowercase			  =							JsonDatasetReader(lowerCamelCase_						,    cache_dir=lowerCamelCase_			).read()
			_check_json_dataset(lowerCamelCase_						,    lowerCamelCase_			)




def 				_lowerCAmelCase  (			lowerCamelCase_     :					str						,    lowerCamelCase_     :					str						,    lowerCamelCase_     :					Tuple=("train",)			):
			assert isinstance(lowerCamelCase_						,    lowerCamelCase_			)
			for split in splits:
						__lowercase			  =							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 				_lowerCAmelCase  (			lowerCamelCase_     :					List[str]						,    lowerCamelCase_     :					Any						,    lowerCamelCase_     :					Dict			):
			__lowercase			  =							tmp_path / '''cache'''
			__lowercase			  =							{'''col_1''': '''string''', '''col_2''': '''int64''', '''col_3''': '''float64'''}
			with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
						__lowercase			  =							JsonDatasetReader({'''train''': jsonl_path}						,    cache_dir=lowerCamelCase_						,    keep_in_memory=lowerCamelCase_			).read()
			_check_json_datasetdict(lowerCamelCase_						,    lowerCamelCase_			)




@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 				_lowerCAmelCase  (			lowerCamelCase_     :					Optional[int]						,    lowerCamelCase_     :					Union[str, Any]						,    lowerCamelCase_     :					int			):
			__lowercase			  =							tmp_path / '''cache'''
			__lowercase			  =							{'''col_1''': '''string''', '''col_2''': '''int64''', '''col_3''': '''float64'''}
			__lowercase			  =							features.copy() if features else default_expected_features
			__lowercase			  =							(
			    Features({feature: Value(lowerCamelCase_			) for feature, dtype in features.items()}			) if features is not None else None
			)
			__lowercase			  =							JsonDatasetReader({'''train''': jsonl_path}						,    features=lowerCamelCase_						,    cache_dir=lowerCamelCase_			).read()
			_check_json_datasetdict(lowerCamelCase_						,    lowerCamelCase_			)




@pytest.mark.parametrize('''split'''						,    [None, NamedSplit('''train'''			), '''train''', '''test''']			)
def 				_lowerCAmelCase  (			lowerCamelCase_     :					Union[str, Any]						,    lowerCamelCase_     :					List[str]						,    lowerCamelCase_     :					int			):
			if split:
						__lowercase			  =							{split: jsonl_path}
			else:
						__lowercase			  =							'''train'''
						__lowercase			  =							{'''train''': jsonl_path, '''test''': jsonl_path}
			__lowercase			  =							tmp_path / '''cache'''
			__lowercase			  =							{'''col_1''': '''string''', '''col_2''': '''int64''', '''col_3''': '''float64'''}
			__lowercase			  =							JsonDatasetReader(lowerCamelCase_						,    cache_dir=lowerCamelCase_			).read()
			_check_json_datasetdict(lowerCamelCase_						,    lowerCamelCase_						,    splits=list(path.keys()			)			)
			assert all(dataset[split].split == split for split in path.keys()			)




def 				_lowerCAmelCase  (			lowerCamelCase_     :					List[str]			):
			return json.load(lowerCamelCase_			)




def 				_lowerCAmelCase  (			lowerCamelCase_     :					List[Any]			):
			return [json.loads(lowerCamelCase_			) for line in buffer]






class 							__lowercase    :







	'''simple docstring'''

	@pytest.mark.parametrize('''lines, load_json_function'''       ,[(True, load_json_lines), (False, load_json)]	)
	def   _UpperCAmelCase       (self       ,_lowerCamelCase       ,_lowerCamelCase       ,_lowerCamelCase	)							->							Optional[Any]:







				'''simple docstring'''





				with io.BytesIO() as buffer:
							JsonDatasetWriter(_lowerCamelCase       ,_lowerCamelCase       ,lines=_lowerCamelCase	).write()
							buffer.seek(0	)
							__lowercase			  =							load_json_function(_lowerCamelCase	)
				assert isinstance(_lowerCamelCase       ,_lowerCamelCase	)
				assert isinstance(exported_content[0]       ,_lowerCamelCase	)
				assert len(_lowerCamelCase	) == 10


	@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   _UpperCAmelCase       (self       ,_lowerCamelCase       ,_lowerCamelCase       ,_lowerCamelCase       ,_lowerCamelCase       ,_lowerCamelCase	)							->							Union[str, Any]:







				'''simple docstring'''





				with io.BytesIO() as buffer:
							JsonDatasetWriter(_lowerCamelCase       ,_lowerCamelCase       ,lines=_lowerCamelCase       ,orient=_lowerCamelCase	).write()
							buffer.seek(0	)
							__lowercase			  =							load_json(_lowerCamelCase	)
				assert isinstance(_lowerCamelCase       ,_lowerCamelCase	)
				if keys:
							if container is dict:
										assert exported_content.keys() == keys
							else:
										assert exported_content[0].keys() == keys
				else:
							assert not hasattr(_lowerCamelCase       ,'''keys'''	) and not hasattr(exported_content[0]       ,'''keys'''	)
				if len_at:
							assert len(exported_content[len_at]	) == 10
				else:
							assert len(_lowerCamelCase	) == 10


	@pytest.mark.parametrize('''lines, load_json_function'''       ,[(True, load_json_lines), (False, load_json)]	)
	def   _UpperCAmelCase       (self       ,_lowerCamelCase       ,_lowerCamelCase       ,_lowerCamelCase	)							->							Union[str, Any]:







				'''simple docstring'''





				with io.BytesIO() as buffer:
							JsonDatasetWriter(_lowerCamelCase       ,_lowerCamelCase       ,lines=_lowerCamelCase       ,num_proc=2	).write()
							buffer.seek(0	)
							__lowercase			  =							load_json_function(_lowerCamelCase	)
				assert isinstance(_lowerCamelCase       ,_lowerCamelCase	)
				assert isinstance(exported_content[0]       ,_lowerCamelCase	)
				assert len(_lowerCamelCase	) == 10


	@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   _UpperCAmelCase       (self       ,_lowerCamelCase       ,_lowerCamelCase       ,_lowerCamelCase       ,_lowerCamelCase       ,_lowerCamelCase	)							->							List[str]:







				'''simple docstring'''





				with io.BytesIO() as buffer:
							JsonDatasetWriter(_lowerCamelCase       ,_lowerCamelCase       ,lines=_lowerCamelCase       ,orient=_lowerCamelCase       ,num_proc=2	).write()
							buffer.seek(0	)
							__lowercase			  =							load_json(_lowerCamelCase	)
				assert isinstance(_lowerCamelCase       ,_lowerCamelCase	)
				if keys:
							if container is dict:
										assert exported_content.keys() == keys
							else:
										assert exported_content[0].keys() == keys
				else:
							assert not hasattr(_lowerCamelCase       ,'''keys'''	) and not hasattr(exported_content[0]       ,'''keys'''	)
				if len_at:
							assert len(exported_content[len_at]	) == 10
				else:
							assert len(_lowerCamelCase	) == 10


	def   _UpperCAmelCase       (self       ,_lowerCamelCase	)							->							Dict:







				'''simple docstring'''





				with pytest.raises(_lowerCamelCase	):
							with io.BytesIO() as buffer:
										JsonDatasetWriter(_lowerCamelCase       ,_lowerCamelCase       ,num_proc=0	)


	@pytest.mark.parametrize('''compression, extension'''       ,[('''gzip''', '''gz'''), ('''bz2''', '''bz2'''), ('''xz''', '''xz''')]	)
	def   _UpperCAmelCase       (self       ,_lowerCamelCase       ,_lowerCamelCase       ,_lowerCamelCase       ,_lowerCamelCase       ,_lowerCamelCase	)							->							Optional[int]:







				'''simple docstring'''





				__lowercase			  =							tmp_path_factory.mktemp('''data'''	) / f"test.json.{extension}"
				__lowercase			  =							str(shared_datadir / f"test_file.json.{extension}"	)
				JsonDatasetWriter(_lowerCamelCase       ,_lowerCamelCase       ,compression=_lowerCamelCase	).write()

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












'''simple docstring'''



def 				_lowerCAmelCase  (			lowerCamelCase_     :					Union[str, Any]			):
			__lowercase			  =							1
			__lowercase			  =							2
			while i * i <= n:
						__lowercase			  =							0
						while n % i == 0:
									n //= i
									multiplicity += 1
						n_divisors *= multiplicity + 1
						i += 1
			if n > 1:
						n_divisors *= 2
			return n_divisors




def 				_lowerCAmelCase  (			):
			__lowercase			  =							1
			__lowercase			  =							1

			while True:
						i += 1
						t_num += i

						if count_divisors(lowerCamelCase_			) > 5_0_0:
									break

			return t_num


if __name__ == "__main__":
	print(solution())





