


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__			(      a__:		str      ,  a__:		List[Any]	)  ->		str:


							'''simple docstring'''





							assert isinstance(a__      ,  a__	)
							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__			(      a__:		Tuple      ,  a__:		Dict      ,  a__:		Any	)  ->		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(a__      ,  cache_dir=a__      ,  keep_in_memory=a__	).read()
							_check_json_dataset(a__      ,  a__	)

@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__			(      a__:		Any      ,  a__:		int      ,  a__:		Dict	)  ->		Optional[Any]:


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

@pytest.mark.parametrize(
    'features'      ,  [
        None,
        {'col_3': 'float64', 'col_1': 'string', 'col_2': 'int64'},
    ]      ,  )
def 	lowerCAmelCase__			(      a__:		Any      ,  a__:		Tuple      ,  a__:		int	)  ->		Any:


							'''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(a__	) for feature, dtype in features.items()}	) if features is not None else None
							)
							_UpperCAmelCase					=      JsonDatasetReader(a__      ,  features=a__      ,  cache_dir=a__	).read()
							assert isinstance(a__      ,  a__	)
							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__			(      a__:		Dict      ,  a__:		List[str]	)  ->		List[str]:


							'''simple docstring'''





							_UpperCAmelCase					=      {'col_2': 'int64', 'col_3': 'float64', 'col_1': 'string'}
							_UpperCAmelCase					=      features.copy()
							_UpperCAmelCase					=      (
							    Features({feature: Value(a__	) for feature, dtype in features.items()}	) if features is not None else None
							)
							_UpperCAmelCase					=      tmp_path / 'cache'
							_UpperCAmelCase					=      JsonDatasetReader(a__      ,  features=a__      ,  cache_dir=a__	).read()
							assert isinstance(a__      ,  a__	)
							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__			(      a__:		Optional[int]      ,  a__:		List[str]      ,  a__:		int	)  ->		int:


							'''simple docstring'''





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

@pytest.mark.parametrize('path_type'      ,  [str, list]	)
def 	lowerCAmelCase__			(      a__:		List[Any]      ,  a__:		Union[str, Any]      ,  a__:		Union[str, Any]	)  ->		Tuple:


							'''simple docstring'''





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

def 	lowerCAmelCase__			(      a__:		List[Any]      ,  a__:		Union[str, Any]      ,  a__:		Union[str, Any]=("train",)	)  ->		int:


							'''simple docstring'''





							assert isinstance(a__      ,  a__	)
							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 	lowerCAmelCase__			(      a__:		Dict      ,  a__:		Union[str, Any]      ,  a__:		int	)  ->		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=a__      ,  keep_in_memory=a__	).read()
							_check_json_datasetdict(a__      ,  a__	)

@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__			(      a__:		List[str]      ,  a__:		int      ,  a__:		List[str]	)  ->		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(a__	) for feature, dtype in features.items()}	) if features is not None else None
							)
							_UpperCAmelCase					=      JsonDatasetReader({'train': jsonl_path}      ,  features=a__      ,  cache_dir=a__	).read()
							_check_json_datasetdict(a__      ,  a__	)

@pytest.mark.parametrize('split'      ,  [None, NamedSplit('train'	), 'train', 'test']	)
def 	lowerCAmelCase__			(      a__:		Tuple      ,  a__:		Dict      ,  a__:		Union[str, Any]	)  ->		List[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(a__      ,  cache_dir=a__	).read()
							_check_json_datasetdict(a__      ,  a__      ,  splits=list(path.keys()	)	)
							assert all(dataset[split].split == split for split in path.keys()	)

def 	lowerCAmelCase__			(      a__:		Tuple	)  ->		Any:


							'''simple docstring'''





							return json.load(a__	)

def 	lowerCAmelCase__			(      a__:		Optional[Any]	)  ->		Optional[Any]:


							'''simple docstring'''





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







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




									"""simple docstring"""





									with io.BytesIO() as buffer:
																JsonDatasetWriter(_SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE      ,   lines=_SCREAMING_SNAKE_CASE						).write()
																buffer.seek(0						)
																_UpperCAmelCase					=      load_json_function(_SCREAMING_SNAKE_CASE						)
									assert isinstance(_SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE						)
									assert isinstance(exported_content[0]      ,   _SCREAMING_SNAKE_CASE						)
									assert len(_SCREAMING_SNAKE_CASE						) == 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      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE						)			->	int:




									"""simple docstring"""





									with io.BytesIO() as buffer:
																JsonDatasetWriter(_SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE      ,   lines=_SCREAMING_SNAKE_CASE      ,   orient=_SCREAMING_SNAKE_CASE						).write()
																buffer.seek(0						)
																_UpperCAmelCase					=      load_json(_SCREAMING_SNAKE_CASE						)
									assert isinstance(_SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE						)
									if keys:
																if container is dict:
																							assert exported_content.keys() == keys
																else:
																							assert exported_content[0].keys() == keys
									else:
																assert not hasattr(_SCREAMING_SNAKE_CASE      ,   'keys'						) and not hasattr(exported_content[0]      ,   'keys'						)
									if len_at:
																assert len(exported_content[len_at]						) == 10
									else:
																assert len(_SCREAMING_SNAKE_CASE						) == 10
		@pytest.mark.parametrize('lines, load_json_function'      ,   [(True, load_json_lines), (False, load_json)]						)
		def      UpperCAmelCase__    (      self      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE						)			->	Union[str, Any]:




									"""simple docstring"""





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




									"""simple docstring"""





									with io.BytesIO() as buffer:
																JsonDatasetWriter(_SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE      ,   lines=_SCREAMING_SNAKE_CASE      ,   orient=_SCREAMING_SNAKE_CASE      ,   num_proc=2						).write()
																buffer.seek(0						)
																_UpperCAmelCase					=      load_json(_SCREAMING_SNAKE_CASE						)
									assert isinstance(_SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE						)
									if keys:
																if container is dict:
																							assert exported_content.keys() == keys
																else:
																							assert exported_content[0].keys() == keys
									else:
																assert not hasattr(_SCREAMING_SNAKE_CASE      ,   'keys'						) and not hasattr(exported_content[0]      ,   'keys'						)
									if len_at:
																assert len(exported_content[len_at]						) == 10
									else:
																assert len(_SCREAMING_SNAKE_CASE						) == 10
		def      UpperCAmelCase__    (      self      ,   _SCREAMING_SNAKE_CASE						)			->	Optional[int]:




									"""simple docstring"""





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






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




									"""simple docstring"""





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

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




from importlib import import_module

from .logging import get_logger


lowerCAmelCase__  :Optional[Any]             =							get_logger(__name__)

class 		__a				:
		def __init__(      self      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE=None						)			->	List[Any]:




									"""simple docstring"""





									_UpperCAmelCase					=      attrs or []
									if module is not None:
																for key in module.__dict__:
																							if key in attrs or not key.startswith('__'						):
																														setattr(self      ,   _SCREAMING_SNAKE_CASE      ,   getattr(_SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE						)						)
									_UpperCAmelCase					=      module._original_module if isinstance(_SCREAMING_SNAKE_CASE      ,   _PatchedModuleObj						) else module



class 		__a				:
		_a     :  Any           =     []
		def __init__(      self      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE=None						)			->	Union[str, Any]:




									"""simple docstring"""





									_UpperCAmelCase					=      obj
									_UpperCAmelCase					=      target
									_UpperCAmelCase					=      new
									_UpperCAmelCase					=      target.split('.'						)[0]
									_UpperCAmelCase					=      {}
									_UpperCAmelCase					=      attrs or []
		def __enter__(      self						)			->	int:




									"""simple docstring"""





									*_UpperCAmelCase		,	_UpperCAmelCase					=      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(_SCREAMING_SNAKE_CASE						)						):
																try:
																							_UpperCAmelCase					=      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__():
																							_UpperCAmelCase					=      getattr(self.obj      ,   _SCREAMING_SNAKE_CASE						)
																							# 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(_SCREAMING_SNAKE_CASE      ,   _PatchedModuleObj						) and obj_attr._original_module is submodule)
																							):
																														_UpperCAmelCase					=      obj_attr
																														# patch at top level
																														setattr(self.obj      ,   _SCREAMING_SNAKE_CASE      ,   _PatchedModuleObj(_SCREAMING_SNAKE_CASE      ,   attrs=self.attrs						)						)
																														_UpperCAmelCase					=      getattr(self.obj      ,   _SCREAMING_SNAKE_CASE						)
																														# construct lower levels patches
																														for key in submodules[i + 1 :]:
																																					setattr(_SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE      ,   _PatchedModuleObj(getattr(_SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE						)      ,   attrs=self.attrs						)						)
																																					_UpperCAmelCase					=      getattr(_SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE						)
																														# finally set the target attribute
																														setattr(_SCREAMING_SNAKE_CASE      ,   _SCREAMING_SNAKE_CASE      ,   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:
																							_UpperCAmelCase					=      getattr(import_module('.'.join(_SCREAMING_SNAKE_CASE						)						)      ,   _SCREAMING_SNAKE_CASE						)
																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      ,   _SCREAMING_SNAKE_CASE						) is attr_value:
																														_UpperCAmelCase					=      getattr(self.obj      ,   _SCREAMING_SNAKE_CASE						)
																														setattr(self.obj      ,   _SCREAMING_SNAKE_CASE      ,   self.new						)
									elif target_attr in globals()["__builtins__"]:  # if it'a s builtin like "open"
																_UpperCAmelCase					=      globals()['__builtins__'][target_attr]
																setattr(self.obj      ,   _SCREAMING_SNAKE_CASE      ,   self.new						)
									else:
																raise RuntimeError(f'''Tried to patch attribute {target_attr} instead of a submodule.'''						)
		def __exit__(      self      ,   *_SCREAMING_SNAKE_CASE						)			->	Union[str, Any]:




									"""simple docstring"""





									for attr in list(self.original						):
																setattr(self.obj      ,   _SCREAMING_SNAKE_CASE      ,   self.original.pop(_SCREAMING_SNAKE_CASE						)						)
		def      UpperCAmelCase__    (      self						)			->	Optional[int]:




									"""simple docstring"""





									self.__enter__()
									self._active_patches.append(self						)






		def      UpperCAmelCase__    (      self						)			->	str:




									"""simple docstring"""





									try:
																self._active_patches.remove(self						)
									except ValueError:
																# If the patch hasn't been started this will fail
																return None

									return self.__exit__()
