


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


			'''simple docstring'''







			assert isinstance(_UpperCamelCase							,      _UpperCamelCase	)
			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_	(  _UpperCamelCase							,      _UpperCamelCase							,      _UpperCamelCase	):


			'''simple docstring'''







			__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(_UpperCamelCase							,      cache_dir=_UpperCamelCase							,      keep_in_memory=_UpperCamelCase	).read()
			_check_json_dataset(_UpperCamelCase							,      _UpperCamelCase	)


@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_	(  _UpperCamelCase							,      _UpperCamelCase							,      _UpperCamelCase	):


			'''simple docstring'''







			__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(_UpperCamelCase	) for feature, dtype in features.items()}	) if features is not None else None
			)
			__lowercase			       = JsonDatasetReader(_UpperCamelCase							,      features=_UpperCamelCase							,      cache_dir=_UpperCamelCase	).read()
			_check_json_dataset(_UpperCamelCase							,      _UpperCamelCase	)


@pytest.mark.parametrize(
    '''features'''							,      [
        None,
        {'''col_3''': '''float64''', '''col_1''': '''string''', '''col_2''': '''int64'''},
    ]							,      )
def 		lowercase_	(  _UpperCamelCase							,      _UpperCamelCase							,      _UpperCamelCase	):


			'''simple docstring'''







			__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(_UpperCamelCase	) for feature, dtype in features.items()}	) if features is not None else None
			)
			__lowercase			       = JsonDatasetReader(_UpperCamelCase							,      features=_UpperCamelCase							,      cache_dir=_UpperCamelCase	).read()
			assert isinstance(_UpperCamelCase							,      _UpperCamelCase	)
			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_	(  _UpperCamelCase							,      _UpperCamelCase	):


			'''simple docstring'''







			__lowercase			       = {'''col_2''': '''int64''', '''col_3''': '''float64''', '''col_1''': '''string'''}
			__lowercase			       = features.copy()
			__lowercase			       = (
			    Features({feature: Value(_UpperCamelCase	) for feature, dtype in features.items()}	) if features is not None else None
			)
			__lowercase			       = tmp_path / '''cache'''
			__lowercase			       = JsonDatasetReader(_UpperCamelCase							,      features=_UpperCamelCase							,      cache_dir=_UpperCamelCase	).read()
			assert isinstance(_UpperCamelCase							,      _UpperCamelCase	)
			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_	(  _UpperCamelCase							,      _UpperCamelCase							,      _UpperCamelCase	):


			'''simple docstring'''







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


@pytest.mark.parametrize('''path_type'''							,      [str, list]	)
def 		lowercase_	(  _UpperCamelCase							,      _UpperCamelCase							,      _UpperCamelCase	):


			'''simple docstring'''







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


def 		lowercase_	(  _UpperCamelCase							,      _UpperCamelCase							,      _UpperCamelCase=("train",)	):


			'''simple docstring'''







			assert isinstance(_UpperCamelCase							,      _UpperCamelCase	)
			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 		lowercase_	(  _UpperCamelCase							,      _UpperCamelCase							,      _UpperCamelCase	):


			'''simple docstring'''







			__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=_UpperCamelCase							,      keep_in_memory=_UpperCamelCase	).read()
			_check_json_datasetdict(_UpperCamelCase							,      _UpperCamelCase	)


@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_	(  _UpperCamelCase							,      _UpperCamelCase							,      _UpperCamelCase	):


			'''simple docstring'''







			__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(_UpperCamelCase	) for feature, dtype in features.items()}	) if features is not None else None
			)
			__lowercase			       = JsonDatasetReader({'''train''': jsonl_path}							,      features=_UpperCamelCase							,      cache_dir=_UpperCamelCase	).read()
			_check_json_datasetdict(_UpperCamelCase							,      _UpperCamelCase	)


@pytest.mark.parametrize('''split'''							,      [None, NamedSplit('''train'''	), '''train''', '''test''']	)
def 		lowercase_	(  _UpperCamelCase							,      _UpperCamelCase							,      _UpperCamelCase	):


			'''simple docstring'''







			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(_UpperCamelCase							,      cache_dir=_UpperCamelCase	).read()
			_check_json_datasetdict(_UpperCamelCase							,      _UpperCamelCase							,      splits=list(path.keys()	)	)
			assert all(dataset[split].split == split for split in path.keys()	)


def 		lowercase_	(  _UpperCamelCase	):


			'''simple docstring'''







			return json.load(_UpperCamelCase	)


def 		lowercase_	(  _UpperCamelCase	):


			'''simple docstring'''







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





class 						lowerCamelCase_							:



					'''simple docstring'''







					@pytest.mark.parametrize('''lines, load_json_function'''			,							[(True, load_json_lines), (False, load_json)]    )
					def A   (					self			,							snake_case_			,							snake_case_			,							snake_case_    ) ->      List[Any]:







								'''simple docstring'''

								with io.BytesIO() as buffer:
											JsonDatasetWriter(snake_case_			,							snake_case_			,							lines=snake_case_    ).write()
											buffer.seek(0    )
											__lowercase			       = load_json_function(snake_case_    )
								assert isinstance(snake_case_			,							snake_case_    )
								assert isinstance(exported_content[0]			,							snake_case_    )
								assert len(snake_case_    ) == 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 A   (					self			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_    ) ->      List[Any]:







								'''simple docstring'''

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






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







								'''simple docstring'''

								with io.BytesIO() as buffer:
											JsonDatasetWriter(snake_case_			,							snake_case_			,							lines=snake_case_			,							num_proc=2    ).write()
											buffer.seek(0    )
											__lowercase			       = load_json_function(snake_case_    )
								assert isinstance(snake_case_			,							snake_case_    )
								assert isinstance(exported_content[0]			,							snake_case_    )
								assert len(snake_case_    ) == 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 A   (					self			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_    ) ->      Tuple:







								'''simple docstring'''

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






					def A   (					self			,							snake_case_    ) ->      List[str]:







								'''simple docstring'''

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






					@pytest.mark.parametrize('''compression, extension'''			,							[('''gzip''', '''gz'''), ('''bz2''', '''bz2'''), ('''xz''', '''xz''')]    )
					def A   (					self			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_    ) ->      str:







								'''simple docstring'''

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

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






from typing import Optional

from .. import Features, NamedSplit
from ..packaged_modules.text.text import Text
from ..utils.typing import NestedDataStructureLike, PathLike
from .abc import AbstractDatasetReader
class 						lowerCamelCase_							(					lowerCAmelCase__					):



					'''simple docstring'''







					def __init__(					self			,							snake_case_			,							snake_case_ = None			,							snake_case_ = None			,							snake_case_ = None			,							snake_case_ = False			,							snake_case_ = False			,							snake_case_ = None			,							**snake_case_			,							) ->      Dict:







								'''simple docstring'''

								super().__init__(
								    snake_case_			,							split=snake_case_			,							features=snake_case_			,							cache_dir=snake_case_			,							keep_in_memory=snake_case_			,							streaming=snake_case_			,							num_proc=snake_case_			,							**snake_case_			,							)
								__lowercase			       = path_or_paths if isinstance(snake_case_			,							snake_case_    ) else {self.split: path_or_paths}
								__lowercase			       = Text(
								    cache_dir=snake_case_			,							data_files=snake_case_			,							features=snake_case_			,							**snake_case_			,							)






					def A   (					self    ) ->      int:







								'''simple docstring'''

								if self.streaming:
											__lowercase			       = self.builder.as_streaming_dataset(split=self.split    )
								# Build regular (map-style) dataset
								else:
											__lowercase			       = None
											__lowercase			       = None
											__lowercase			       = None
											__lowercase			       = None

											self.builder.download_and_prepare(
											    download_config=snake_case_			,							download_mode=snake_case_			,							verification_mode=snake_case_			,							base_path=snake_case_			,							num_proc=self.num_proc			,							)
											__lowercase			       = self.builder.as_dataset(
											    split=self.split			,							verification_mode=snake_case_			,							in_memory=self.keep_in_memory    )
								return dataset


