






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

						'''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 _lowerCAmelCase	(  __lowerCamelCase:Optional[Any] ,						__lowerCamelCase:Optional[Any] ,						__lowerCamelCase:Optional[int]   ):

						'''simple docstring'''






						__magic_name__            =     tmp_path / "cache"
						__magic_name__            =     {"col_1": "string", "col_2": "int64", "col_3": "float64"}
						with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
												__magic_name__            =     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 _lowerCAmelCase	(  __lowerCamelCase:Optional[Any] ,						__lowerCamelCase:Optional[Any] ,						__lowerCamelCase:Any   ):

						'''simple docstring'''






						__magic_name__            =     tmp_path / "cache"
						__magic_name__            =     {"col_1": "string", "col_2": "int64", "col_3": "float64"}
						__magic_name__            =     features.copy() if features else default_expected_features
						__magic_name__            =     (
						    Features({feature: Value(_SCREAMING_SNAKE_CASE   ) for feature, dtype in features.items()}   ) if features is not None else None
						)
						__magic_name__            =     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 _lowerCAmelCase	(  __lowerCamelCase:List[str] ,						__lowerCamelCase:List[Any] ,						__lowerCamelCase:Dict   ):

						'''simple docstring'''






						__magic_name__            =     tmp_path / "cache"
						__magic_name__            =     {"col_3": "float64", "col_1": "string", "col_2": "int64"}
						__magic_name__            =     features.copy() if features else default_expected_features
						__magic_name__            =     (
						    Features({feature: Value(_SCREAMING_SNAKE_CASE   ) for feature, dtype in features.items()}   ) if features is not None else None
						)
						__magic_name__            =     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 _lowerCAmelCase	(  __lowerCamelCase:Tuple ,						__lowerCamelCase:List[str]   ):

						'''simple docstring'''






						__magic_name__            =     {"col_2": "int64", "col_3": "float64", "col_1": "string"}
						__magic_name__            =     features.copy()
						__magic_name__            =     (
						    Features({feature: Value(_SCREAMING_SNAKE_CASE   ) for feature, dtype in features.items()}   ) if features is not None else None
						)
						__magic_name__            =     tmp_path / "cache"
						__magic_name__            =     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 _lowerCAmelCase	(  __lowerCamelCase:List[str] ,						__lowerCamelCase:int ,						__lowerCamelCase:Any   ):

						'''simple docstring'''






						__magic_name__            =     tmp_path / "cache"
						__magic_name__            =     {"col_1": "string", "col_2": "int64", "col_3": "float64"}
						__magic_name__            =     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 _lowerCAmelCase	(  __lowerCamelCase:Optional[int] ,						__lowerCamelCase:Any ,						__lowerCamelCase:Any   ):

						'''simple docstring'''






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






def _lowerCAmelCase	(  __lowerCamelCase:Union[str, Any] ,						__lowerCamelCase:Tuple ,						__lowerCamelCase:str=("train",)   ):

						'''simple docstring'''






						assert isinstance(_SCREAMING_SNAKE_CASE ,						_SCREAMING_SNAKE_CASE   )
						for split in splits:
												__magic_name__            =     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:Optional[int] ,						__lowerCamelCase:Union[str, Any]   ):

						'''simple docstring'''






						__magic_name__            =     tmp_path / "cache"
						__magic_name__            =     {"col_1": "string", "col_2": "int64", "col_3": "float64"}
						with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
												__magic_name__            =     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 _lowerCAmelCase	(  __lowerCamelCase:Any ,						__lowerCamelCase:int ,						__lowerCamelCase:Tuple   ):

						'''simple docstring'''






						__magic_name__            =     tmp_path / "cache"
						__magic_name__            =     {"col_1": "string", "col_2": "int64", "col_3": "float64"}
						__magic_name__            =     features.copy() if features else default_expected_features
						__magic_name__            =     (
						    Features({feature: Value(_SCREAMING_SNAKE_CASE   ) for feature, dtype in features.items()}   ) if features is not None else None
						)
						__magic_name__            =     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 _lowerCAmelCase	(  __lowerCamelCase:Any ,						__lowerCamelCase:Optional[int] ,						__lowerCamelCase:List[Any]   ):

						'''simple docstring'''






						if split:
												__magic_name__            =     {split: jsonl_path}
						else:
												__magic_name__            =     "train"
												__magic_name__            =     {"train": jsonl_path, "test": jsonl_path}
						__magic_name__            =     tmp_path / "cache"
						__magic_name__            =     {"col_1": "string", "col_2": "int64", "col_3": "float64"}
						__magic_name__            =     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 _lowerCAmelCase	(  __lowerCamelCase:List[Any]   ):

						'''simple docstring'''






						return json.load(_SCREAMING_SNAKE_CASE   )






def _lowerCAmelCase	(  __lowerCamelCase:int   ):

						'''simple docstring'''






						return [json.loads(_SCREAMING_SNAKE_CASE   ) for line in buffer]


class 						A_							:







				@pytest.mark.parametrize("lines, load_json_function"						,    [(True, load_json_lines), (False, load_json)]	)
				def     _snake_case					(							self							:						Any						,    __lowerCamelCase							:						int						,    __lowerCamelCase							:						Optional[int]						,    __lowerCamelCase							:						Tuple	)				->					Optional[int]:







										with io.BytesIO() as buffer:
																JsonDatasetWriter(__lowerCamelCase						,    __lowerCamelCase						,    lines=__lowerCamelCase	).write()
																buffer.seek(0	)
																__magic_name__            =     load_json_function(__lowerCamelCase	)
										assert isinstance(__lowerCamelCase						,    __lowerCamelCase	)
										assert isinstance(exported_content[0]						,    __lowerCamelCase	)
										assert len(__lowerCamelCase	) == 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     _snake_case					(							self							:						int						,    __lowerCamelCase							:						Optional[int]						,    __lowerCamelCase							:						str						,    __lowerCamelCase							:						Dict						,    __lowerCamelCase							:						Any						,    __lowerCamelCase							:						Dict	)				->					List[str]:







										with io.BytesIO() as buffer:
																JsonDatasetWriter(__lowerCamelCase						,    __lowerCamelCase						,    lines=__lowerCamelCase						,    orient=__lowerCamelCase	).write()
																buffer.seek(0	)
																__magic_name__            =     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]	) == 1_0
										else:
																assert len(__lowerCamelCase	) == 1_0







				@pytest.mark.parametrize("lines, load_json_function"						,    [(True, load_json_lines), (False, load_json)]	)
				def     _snake_case					(							self							:						Tuple						,    __lowerCamelCase							:						List[Any]						,    __lowerCamelCase							:						List[Any]						,    __lowerCamelCase							:						str	)				->					Any:







										with io.BytesIO() as buffer:
																JsonDatasetWriter(__lowerCamelCase						,    __lowerCamelCase						,    lines=__lowerCamelCase						,    num_proc=2	).write()
																buffer.seek(0	)
																__magic_name__            =     load_json_function(__lowerCamelCase	)
										assert isinstance(__lowerCamelCase						,    __lowerCamelCase	)
										assert isinstance(exported_content[0]						,    __lowerCamelCase	)
										assert len(__lowerCamelCase	) == 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     _snake_case					(							self							:						Optional[Any]						,    __lowerCamelCase							:						Tuple						,    __lowerCamelCase							:						Dict						,    __lowerCamelCase							:						Any						,    __lowerCamelCase							:						str						,    __lowerCamelCase							:						str	)				->					List[str]:







										with io.BytesIO() as buffer:
																JsonDatasetWriter(__lowerCamelCase						,    __lowerCamelCase						,    lines=__lowerCamelCase						,    orient=__lowerCamelCase						,    num_proc=2	).write()
																buffer.seek(0	)
																__magic_name__            =     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]	) == 1_0
										else:
																assert len(__lowerCamelCase	) == 1_0







				def     _snake_case					(							self							:						int						,    __lowerCamelCase							:						int	)				->					List[str]:







										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     _snake_case					(							self							:						str						,    __lowerCamelCase							:						List[str]						,    __lowerCamelCase							:						int						,    __lowerCamelCase							:						List[str]						,    __lowerCamelCase							:						Dict						,    __lowerCamelCase							:						Union[str, Any]	)				->					Dict:
										__magic_name__            =     tmp_path_factory.mktemp("data"	) / f'''test.json.{extension}'''
										__magic_name__            =     str(shared_datadir / f'''test_file.json.{extension}'''	)
										JsonDatasetWriter(__lowerCamelCase						,    __lowerCamelCase						,    compression=__lowerCamelCase	).write()

										with fsspec.open(__lowerCamelCase						,    "rb"						,    compression="infer"	) as f:
																__magic_name__            =     f.read()
										with fsspec.open(__lowerCamelCase						,    "rb"						,    compression="infer"	) as f:
																__magic_name__            =     f.read()
										assert exported_content == original_content











"""simple docstring"""

def _lowerCAmelCase	(  __lowerCamelCase:int   ):

						'''simple docstring'''






						if number < 0:
												raise ValueError("number must not be negative"   )
						return number & (number - 1) == 0


if __name__ == "__main__":
			import doctest

			doctest.testmod()



