"""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 			lowercase__(       __SCREAMING_SNAKE_CASE							:     Any     ,			__SCREAMING_SNAKE_CASE							:     List[Any]      ):
     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 			lowercase__(       __SCREAMING_SNAKE_CASE							:     Tuple     ,			__SCREAMING_SNAKE_CASE							:     str     ,			__SCREAMING_SNAKE_CASE							:     Any      ):
     lowercase_	:				Optional[int]													=	tmp_path / 'cache'
     lowercase_	:				int													=	{'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_	:				Optional[Any]													=	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 			lowercase__(       __SCREAMING_SNAKE_CASE							:     Union[str, Any]     ,			__SCREAMING_SNAKE_CASE							:     Optional[int]     ,			__SCREAMING_SNAKE_CASE							:     Optional[Any]      ):
     lowercase_	:				Union[str, Any]													=	tmp_path / 'cache'
     lowercase_	:				Optional[Any]													=	{'col_1': 'string', 'col_2': 'int64', 'col_3': 'float64'}
     lowercase_	:				List[str]													=	features.copy() if features else default_expected_features
     lowercase_	:				Union[str, Any]													=	(
         Features({feature: Value(__SCREAMING_SNAKE_CASE      ) for feature, dtype in features.items()}      ) if features is not None else None
     )
     lowercase_	:				Dict													=	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 			lowercase__(       __SCREAMING_SNAKE_CASE							:     int     ,			__SCREAMING_SNAKE_CASE							:     Union[str, Any]     ,			__SCREAMING_SNAKE_CASE							:     List[Any]      ):
     lowercase_	:				Dict													=	tmp_path / 'cache'
     lowercase_	:				Union[str, Any]													=	{'col_3': 'float64', 'col_1': 'string', 'col_2': 'int64'}
     lowercase_	:				Dict													=	features.copy() if features else default_expected_features
     lowercase_	:				Union[str, Any]													=	(
         Features({feature: Value(__SCREAMING_SNAKE_CASE      ) for feature, dtype in features.items()}      ) if features is not None else None
     )
     lowercase_	:				Union[str, Any]													=	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 			lowercase__(       __SCREAMING_SNAKE_CASE							:     int     ,			__SCREAMING_SNAKE_CASE							:     Tuple      ):
     # jsonl_312_path features are {"col_3": "float64", "col_1": "string", "col_2": "int64"}
     lowercase_	:				List[str]													=	{'col_2': 'int64', 'col_3': 'float64', 'col_1': 'string'}
     lowercase_	:				Tuple													=	features.copy()
     lowercase_	:				str													=	(
         Features({feature: Value(__SCREAMING_SNAKE_CASE      ) for feature, dtype in features.items()}      ) if features is not None else None
     )
     lowercase_	:				int													=	tmp_path / 'cache'
     lowercase_	:				Optional[Any]													=	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 			lowercase__(       __SCREAMING_SNAKE_CASE							:     Any     ,			__SCREAMING_SNAKE_CASE							:     Any     ,			__SCREAMING_SNAKE_CASE							:     Optional[int]      ):
     lowercase_	:				str													=	tmp_path / 'cache'
     lowercase_	:				List[Any]													=	{'col_1': 'string', 'col_2': 'int64', 'col_3': 'float64'}
     lowercase_	:				Tuple													=	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 			lowercase__(       __SCREAMING_SNAKE_CASE							:     Any     ,			__SCREAMING_SNAKE_CASE							:     Any     ,			__SCREAMING_SNAKE_CASE							:     List[Any]      ):
     if issubclass(__SCREAMING_SNAKE_CASE     ,			__SCREAMING_SNAKE_CASE      ):
          lowercase_	:				Any													=	jsonl_path
     elif issubclass(__SCREAMING_SNAKE_CASE     ,			__SCREAMING_SNAKE_CASE      ):
          lowercase_	:				List[str]													=	[jsonl_path]
     lowercase_	:				Dict													=	tmp_path / 'cache'
     lowercase_	:				Tuple													=	{'col_1': 'string', 'col_2': 'int64', 'col_3': 'float64'}
     lowercase_	:				List[Any]													=	JsonDatasetReader(__SCREAMING_SNAKE_CASE     ,			cache_dir=__SCREAMING_SNAKE_CASE      ).read()
     _check_json_dataset(__SCREAMING_SNAKE_CASE     ,			__SCREAMING_SNAKE_CASE      )



def 			lowercase__(       __SCREAMING_SNAKE_CASE							:     Optional[int]     ,			__SCREAMING_SNAKE_CASE							:     str     ,			__SCREAMING_SNAKE_CASE							:     List[Any]=("train",)      ):
     assert isinstance(__SCREAMING_SNAKE_CASE     ,			__SCREAMING_SNAKE_CASE      )
     for split in splits:
          lowercase_	:				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 			lowercase__(       __SCREAMING_SNAKE_CASE							:     str     ,			__SCREAMING_SNAKE_CASE							:     Optional[Any]     ,			__SCREAMING_SNAKE_CASE							:     Tuple      ):
     lowercase_	:				Optional[int]													=	tmp_path / 'cache'
     lowercase_	:				Optional[int]													=	{'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_	:				int													=	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 			lowercase__(       __SCREAMING_SNAKE_CASE							:     Optional[int]     ,			__SCREAMING_SNAKE_CASE							:     List[str]     ,			__SCREAMING_SNAKE_CASE							:     List[str]      ):
     lowercase_	:				List[str]													=	tmp_path / 'cache'
     lowercase_	:				List[Any]													=	{'col_1': 'string', 'col_2': 'int64', 'col_3': 'float64'}
     lowercase_	:				Union[str, Any]													=	features.copy() if features else default_expected_features
     lowercase_	:				Optional[Any]													=	(
         Features({feature: Value(__SCREAMING_SNAKE_CASE      ) for feature, dtype in features.items()}      ) if features is not None else None
     )
     lowercase_	:				Optional[int]													=	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 			lowercase__(       __SCREAMING_SNAKE_CASE							:     Tuple     ,			__SCREAMING_SNAKE_CASE							:     Optional[Any]     ,			__SCREAMING_SNAKE_CASE							:     Optional[int]      ):
     if split:
          lowercase_	:				List[Any]													=	{split: jsonl_path}
     else:
          lowercase_	:				Any													=	'train'
          lowercase_	:				Optional[int]													=	{'train': jsonl_path, 'test': jsonl_path}
     lowercase_	:				Optional[Any]													=	tmp_path / 'cache'
     lowercase_	:				Optional[int]													=	{'col_1': 'string', 'col_2': 'int64', 'col_3': 'float64'}
     lowercase_	:				int													=	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 			lowercase__(       __SCREAMING_SNAKE_CASE							:     str      ):
     return json.load(__SCREAMING_SNAKE_CASE      )



def 			lowercase__(       __SCREAMING_SNAKE_CASE							:     int      ):
     return [json.loads(__SCREAMING_SNAKE_CASE      ) for line in buffer]







class        UpperCamelCase		:
   @pytest.mark.parametrize('lines, load_json_function'     ,[(True, load_json_lines), (False, load_json)]			)
   def 			_UpperCAmelCase   (			self     ,__UpperCamelCase     ,__UpperCamelCase     ,__UpperCamelCase			)	->      List[str]:


        '''simple docstring'''


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


        '''simple docstring'''


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


        '''simple docstring'''


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


        '''simple docstring'''


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


        '''simple docstring'''


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







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


        '''simple docstring'''


        lowercase_	:				int													=	tmp_path_factory.mktemp('data'			) / f'''test.json.{extension}'''
        lowercase_	:				Optional[Any]													=	str(shared_datadir / f'''test_file.json.{extension}'''			)
        JsonDatasetWriter(__UpperCamelCase     ,__UpperCamelCase     ,compression=__UpperCamelCase			).write()

        with fsspec.open(__UpperCamelCase     ,'rb'     ,compression='infer'			) as f:
             lowercase_	:				List[Any]													=	f.read()
        with fsspec.open(__UpperCamelCase     ,'rb'     ,compression='infer'			) as f:
             lowercase_	:				List[str]													=	f.read()
        assert exported_content == original_content


"""simple docstring"""







import sys

from .dependency_versions_table import deps
from .utils.versions import require_version, require_version_core


# define which module versions we always want to check at run time
# (usually the ones defined in `install_requires` in setup.py)
#
# order specific notes:
# - tqdm must be checked before tokenizers

__SCREAMING_SNAKE_CASE    	="python tqdm regex requests packaging filelock numpy tokenizers".split()
if sys.version_info < (3, 7):
       pkgs_to_check_at_runtime.append("dataclasses")
if sys.version_info < (3, 8):
       pkgs_to_check_at_runtime.append("importlib_metadata")

for pkg in pkgs_to_check_at_runtime:
       if pkg in deps:
              if pkg == "tokenizers":
                     # must be loaded here, or else tqdm check may fail
                     from .utils import is_tokenizers_available

                     if not is_tokenizers_available():
                            continue  # not required, check version only if installed

              require_version_core(deps[pkg])
       else:
              raise ValueError(F"can't find {pkg} in {deps.keys()}, check dependency_versions_table.py")



def 			lowercase__(       __SCREAMING_SNAKE_CASE							:     int     ,			__SCREAMING_SNAKE_CASE							:     Union[str, Any]=None      ):
     require_version(deps[pkg]     ,			__SCREAMING_SNAKE_CASE      )

