


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


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


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


@pytest.mark.parametrize('''path_type'''    ,       [str, list]      )
def        _UpperCamelCase							(      lowercase__    ,       lowercase__    ,       lowercase__      ):
						if issubclass(lowercase__    ,       lowercase__      ):
												__SCREAMING_SNAKE_CASE						:     Dict					=		jsonl_path
						elif issubclass(lowercase__    ,       lowercase__      ):
												__SCREAMING_SNAKE_CASE						:     int					=		[jsonl_path]
						__SCREAMING_SNAKE_CASE						:     int					=		tmp_path / '''cache'''
						__SCREAMING_SNAKE_CASE						:     Any					=		{'''col_1''': '''string''', '''col_2''': '''int64''', '''col_3''': '''float64'''}
						__SCREAMING_SNAKE_CASE						:     List[Any]					=		JsonDatasetReader(lowercase__    ,       cache_dir=lowercase__      ).read()
						_check_json_dataset(lowercase__    ,       lowercase__      )


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


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


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


def        _UpperCamelCase							(      lowercase__      ):
						return json.load(lowercase__      )


def        _UpperCamelCase							(      lowercase__      ):
						return [json.loads(lowercase__      ) for line in buffer]
class        _lowercase      :






				'''simple docstring'''



				@pytest.mark.parametrize('''lines, load_json_function'''			,			[(True, load_json_lines), (False, load_json)]	)
				def 		__magic_name__(       self			:str			,			lowerCAmelCase__			:Dict			,			lowerCAmelCase__			:List[Any]			,			lowerCAmelCase__			:Any	)		->			Optional[Any]:
										with io.BytesIO() as buffer:
																JsonDatasetWriter(lowerCAmelCase__			,			lowerCAmelCase__			,			lines=lowerCAmelCase__	).write()
																buffer.seek(0	)
																__SCREAMING_SNAKE_CASE						:     Tuple					=		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 		__magic_name__(       self			:Optional[int]			,			lowerCAmelCase__			:List[Any]			,			lowerCAmelCase__			:Union[str, Any]			,			lowerCAmelCase__			:Tuple			,			lowerCAmelCase__			:Union[str, Any]			,			lowerCAmelCase__			:List[Any]	)		->			Optional[Any]:
										with io.BytesIO() as buffer:
																JsonDatasetWriter(lowerCAmelCase__			,			lowerCAmelCase__			,			lines=lowerCAmelCase__			,			orient=lowerCAmelCase__	).write()
																buffer.seek(0	)
																__SCREAMING_SNAKE_CASE						:     List[Any]					=		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 		__magic_name__(       self			:Optional[Any]			,			lowerCAmelCase__			:str			,			lowerCAmelCase__			:Any			,			lowerCAmelCase__			:Union[str, Any]	)		->			Dict:
										with io.BytesIO() as buffer:
																JsonDatasetWriter(lowerCAmelCase__			,			lowerCAmelCase__			,			lines=lowerCAmelCase__			,			num_proc=2	).write()
																buffer.seek(0	)
																__SCREAMING_SNAKE_CASE						:     Tuple					=		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 		__magic_name__(       self			:Any			,			lowerCAmelCase__			:List[Any]			,			lowerCAmelCase__			:Union[str, Any]			,			lowerCAmelCase__			:Tuple			,			lowerCAmelCase__			:Optional[Any]			,			lowerCAmelCase__			:Optional[Any]	)		->			Optional[Any]:
										with io.BytesIO() as buffer:
																JsonDatasetWriter(lowerCAmelCase__			,			lowerCAmelCase__			,			lines=lowerCAmelCase__			,			orient=lowerCAmelCase__			,			num_proc=2	).write()
																buffer.seek(0	)
																__SCREAMING_SNAKE_CASE						:     Tuple					=		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 		__magic_name__(       self			:Dict			,			lowerCAmelCase__			:Optional[Any]	)		->			int:
										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 		__magic_name__(       self			:Any			,			lowerCAmelCase__			:Dict			,			lowerCAmelCase__			:int			,			lowerCAmelCase__			:List[str]			,			lowerCAmelCase__			:int			,			lowerCAmelCase__			:Optional[Any]	)		->			int:
										__SCREAMING_SNAKE_CASE						:     Tuple					=		tmp_path_factory.mktemp('''data'''	) / f'''test.json.{extension}'''
										__SCREAMING_SNAKE_CASE						:     Dict					=		str(shared_datadir / f'''test_file.json.{extension}'''	)
										JsonDatasetWriter(lowerCAmelCase__			,			lowerCAmelCase__			,			compression=lowerCAmelCase__	).write()

										with fsspec.open(lowerCAmelCase__			,			'''rb'''			,			compression='''infer'''	) as f:
																__SCREAMING_SNAKE_CASE						:     Dict					=		f.read()
										with fsspec.open(lowerCAmelCase__			,			'''rb'''			,			compression='''infer'''	) as f:
																__SCREAMING_SNAKE_CASE						:     str					=		f.read()
										assert exported_content == original_content




from importlib import import_module

from .logging import get_logger


__lowerCAmelCase	: str  				=get_logger(__name__)
class        _lowercase      :






				'''simple docstring'''



				def __init__(       self			:List[Any]			,			lowerCAmelCase__			:List[Any]			,			lowerCAmelCase__			:str=None	)		->			int:
										__SCREAMING_SNAKE_CASE						:     List[str]					=		attrs or []
										if module is not None:
																for key in module.__dict__:
																						if key in attrs or not key.startswith('''__'''	):
																												setattr(self			,			lowerCAmelCase__			,			getattr(lowerCAmelCase__			,			lowerCAmelCase__	)	)
										__SCREAMING_SNAKE_CASE						:     Optional[Any]					=		module._original_module if isinstance(lowerCAmelCase__			,			_PatchedModuleObj	) else module




class        _lowercase      :






				'''simple docstring'''



				SCREAMING_SNAKE_CASE__				:			Optional[int]						=		[]


				def __init__(       self			:Tuple			,			lowerCAmelCase__			:Optional[Any]			,			lowerCAmelCase__			:str			,			lowerCAmelCase__			:Any			,			lowerCAmelCase__			:Dict=None	)		->			List[Any]:
										__SCREAMING_SNAKE_CASE						:     Optional[int]					=		obj
										__SCREAMING_SNAKE_CASE						:     str					=		target
										__SCREAMING_SNAKE_CASE						:     Dict					=		new
										__SCREAMING_SNAKE_CASE						:     Union[str, Any]					=		target.split('''.'''	)[0]
										__SCREAMING_SNAKE_CASE						:     List[str]					=		{}
										__SCREAMING_SNAKE_CASE						:     Tuple					=		attrs or []


				def __enter__(       self			:int	)		->			Dict:
										*__SCREAMING_SNAKE_CASE			,			__SCREAMING_SNAKE_CASE						:     Optional[Any]					=		self.target.split('''.'''	)

										# Patch modules:
										# it's used to patch attributes of submodules like "os.path.join";
										# in this case we need to patch "os" and "os.path"

										for i in range(len(lowerCAmelCase__	)	):
																try:
																						__SCREAMING_SNAKE_CASE						:     Any					=		import_module('''.'''.join(submodules[: i + 1]	)	)
																except ModuleNotFoundError:
																						continue
																# We iterate over all the globals in self.obj in case we find "os" or "os.path"
																for attr in self.obj.__dir__():
																						__SCREAMING_SNAKE_CASE						:     Union[str, Any]					=		getattr(self.obj			,			lowerCAmelCase__	)
																						# We don't check for the name of the global, but rather if its value *is* "os" or "os.path".
																						# This allows to patch renamed modules like "from os import path as ospath".
																						if obj_attr is submodule or (
																						    (isinstance(lowerCAmelCase__			,			_PatchedModuleObj	) and obj_attr._original_module is submodule)
																						):
																												__SCREAMING_SNAKE_CASE						:     int					=		obj_attr
																												# patch at top level
																												setattr(self.obj			,			lowerCAmelCase__			,			_PatchedModuleObj(lowerCAmelCase__			,			attrs=self.attrs	)	)
																												__SCREAMING_SNAKE_CASE						:     List[str]					=		getattr(self.obj			,			lowerCAmelCase__	)
																												# construct lower levels patches
																												for key in submodules[i + 1 :]:
																																		setattr(lowerCAmelCase__			,			lowerCAmelCase__			,			_PatchedModuleObj(getattr(lowerCAmelCase__			,			lowerCAmelCase__			,			lowerCAmelCase__	)			,			attrs=self.attrs	)	)
																																		__SCREAMING_SNAKE_CASE						:     Tuple					=		getattr(lowerCAmelCase__			,			lowerCAmelCase__	)
																												# finally set the target attribute
																												setattr(lowerCAmelCase__			,			lowerCAmelCase__			,			self.new	)

        # Patch attribute itself:
        # it's used for builtins like "open",
        # and also to patch "os.path.join" we may also need to patch "join"
        # itself if it was imported as "from os.path import join".

										if submodules:  # if it's an attribute of a submodule like "os.path.join"
																try:
																						__SCREAMING_SNAKE_CASE						:     Union[str, Any]					=		getattr(import_module('''.'''.join(lowerCAmelCase__	)	)			,			lowerCAmelCase__	)
																except (AttributeError, ModuleNotFoundError):
																						return
																# We iterate over all the globals in self.obj in case we find "os.path.join"
																for attr in self.obj.__dir__():
																						# We don't check for the name of the global, but rather if its value *is* "os.path.join".
																						# This allows to patch renamed attributes like "from os.path import join as pjoin".
																						if getattr(self.obj			,			lowerCAmelCase__	) is attr_value:
																												__SCREAMING_SNAKE_CASE						:     Any					=		getattr(self.obj			,			lowerCAmelCase__	)
																												setattr(self.obj			,			lowerCAmelCase__			,			self.new	)
										elif target_attr in globals()["__builtins__"]:  # if it'a s builtin like "open"
																__SCREAMING_SNAKE_CASE						:     Union[str, Any]					=		globals()['''__builtins__'''][target_attr]
																setattr(self.obj			,			lowerCAmelCase__			,			self.new	)
										else:
																raise RuntimeError(f'''Tried to patch attribute {target_attr} instead of a submodule.'''	)


				def __exit__(       self			:str			,			*lowerCAmelCase__			:Union[str, Any]	)		->			Optional[int]:
										for attr in list(self.original	):
																setattr(self.obj			,			lowerCAmelCase__			,			self.original.pop(lowerCAmelCase__	)	)


				def 		__magic_name__(       self			:List[Any]	)		->			List[Any]:
										self.__enter__()
										self._active_patches.append(self	)



				def 		__magic_name__(       self			:Optional[int]	)		->			int:
										try:
																self._active_patches.remove(self	)
										except ValueError:
																# If the patch hasn't been started this will fail
																return None

										return self.__exit__()
