


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    (      _a       ,		_a		)					->							Optional[int]:
     '''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     UpperCamelCase    (      _a       ,		_a       ,		_a		)					->							Optional[Any]:
     '''simple docstring'''





     lowercase_       :str     =       tmp_path / '''cache'''
     lowercase_       :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():
          lowercase_       :List[str]     =       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     UpperCamelCase    (      _a       ,		_a       ,		_a		)					->							str:
     '''simple docstring'''





     lowercase_       :Optional[int]     =       tmp_path / '''cache'''
     lowercase_       :List[str]     =       {'''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(_a		) for feature, dtype in features.items()}		) if features is not None else None
     )
     lowercase_       :Tuple     =       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     UpperCamelCase    (      _a       ,		_a       ,		_a		)					->							Union[str, Any]:
     '''simple docstring'''





     lowercase_       :Optional[Any]     =       tmp_path / '''cache'''
     lowercase_       :Tuple     =       {'''col_3''': '''float64''', '''col_1''': '''string''', '''col_2''': '''int64'''}
     lowercase_       :Dict     =       features.copy() if features else default_expected_features
     lowercase_       :str     =       (
         Features({feature: Value(_a		) for feature, dtype in features.items()}		) if features is not None else None
     )
     lowercase_       :List[str]     =       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     UpperCamelCase    (      _a       ,		_a		)					->							List[str]:
     '''simple docstring'''





     lowercase_       :List[str]     =       {'''col_2''': '''int64''', '''col_3''': '''float64''', '''col_1''': '''string'''}
     lowercase_       :Tuple     =       features.copy()
     lowercase_       :List[str]     =       (
         Features({feature: Value(_a		) for feature, dtype in features.items()}		) if features is not None else None
     )
     lowercase_       :List[str]     =       tmp_path / '''cache'''
     lowercase_       :Tuple     =       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     UpperCamelCase    (      _a       ,		_a       ,		_a		)					->							Optional[int]:
     '''simple docstring'''





     lowercase_       :Optional[int]     =       tmp_path / '''cache'''
     lowercase_       :Tuple     =       {'''col_1''': '''string''', '''col_2''': '''int64''', '''col_3''': '''float64'''}
     lowercase_       :Any     =       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     UpperCamelCase    (      _a       ,		_a       ,		_a		)					->							Optional[int]:
     '''simple docstring'''





     if issubclass(_a       ,		_a		):
          lowercase_       :Optional[int]     =       jsonl_path
     elif issubclass(_a       ,		_a		):
          lowercase_       :Tuple     =       [jsonl_path]
     lowercase_       :int     =       tmp_path / '''cache'''
     lowercase_       :List[Any]     =       {'''col_1''': '''string''', '''col_2''': '''int64''', '''col_3''': '''float64'''}
     lowercase_       :Optional[Any]     =       JsonDatasetReader(_a       ,		cache_dir=_a		).read()
     _check_json_dataset(_a       ,		_a		)




def     UpperCamelCase    (      _a       ,		_a       ,		_a=("train",)		)					->							Optional[Any]:
     '''simple docstring'''





     assert isinstance(_a       ,		_a		)
     for split in splits:
          lowercase_       :Optional[Any]     =       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    (      _a       ,		_a       ,		_a		)					->							List[str]:
     '''simple docstring'''





     lowercase_       :int     =       tmp_path / '''cache'''
     lowercase_       :str     =       {'''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[int]     =       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     UpperCamelCase    (      _a       ,		_a       ,		_a		)					->							Optional[int]:
     '''simple docstring'''





     lowercase_       :int     =       tmp_path / '''cache'''
     lowercase_       :Dict     =       {'''col_1''': '''string''', '''col_2''': '''int64''', '''col_3''': '''float64'''}
     lowercase_       :Optional[int]     =       features.copy() if features else default_expected_features
     lowercase_       :Tuple     =       (
         Features({feature: Value(_a		) for feature, dtype in features.items()}		) if features is not None else None
     )
     lowercase_       :Optional[int]     =       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     UpperCamelCase    (      _a       ,		_a       ,		_a		)					->							Tuple:
     '''simple docstring'''





     if split:
          lowercase_       :Any     =       {split: jsonl_path}
     else:
          lowercase_       :List[Any]     =       '''train'''
          lowercase_       :Any     =       {'''train''': jsonl_path, '''test''': jsonl_path}
     lowercase_       :Any     =       tmp_path / '''cache'''
     lowercase_       :Optional[int]     =       {'''col_1''': '''string''', '''col_2''': '''int64''', '''col_3''': '''float64'''}
     lowercase_       :List[Any]     =       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     UpperCamelCase    (      _a		)					->							str:
     '''simple docstring'''





     return json.load(_a		)




def     UpperCamelCase    (      _a		)					->							Any:
     '''simple docstring'''





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






class      UpperCamelCase					:


   '''simple docstring'''







   @pytest.mark.parametrize('''lines, load_json_function'''   ,							[(True, load_json_lines), (False, load_json)]		)
   def 		UpperCamelCase  (     self   ,							UpperCamelCase_   ,							UpperCamelCase_   ,							UpperCamelCase_		):
        with io.BytesIO() as buffer:
             JsonDatasetWriter(UpperCamelCase_   ,							UpperCamelCase_   ,							lines=UpperCamelCase_		).write()
             buffer.seek(0		)
             lowercase_       :Optional[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_		):
        with io.BytesIO() as buffer:
             JsonDatasetWriter(UpperCamelCase_   ,							UpperCamelCase_   ,							lines=UpperCamelCase_   ,							orient=UpperCamelCase_		).write()
             buffer.seek(0		)
             lowercase_       :Optional[int]     =       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_		):
        with io.BytesIO() as buffer:
             JsonDatasetWriter(UpperCamelCase_   ,							UpperCamelCase_   ,							lines=UpperCamelCase_   ,							num_proc=2		).write()
             buffer.seek(0		)
             lowercase_       :int     =       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_		):
        with io.BytesIO() as buffer:
             JsonDatasetWriter(UpperCamelCase_   ,							UpperCamelCase_   ,							lines=UpperCamelCase_   ,							orient=UpperCamelCase_   ,							num_proc=2		).write()
             buffer.seek(0		)
             lowercase_       :Union[str, 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_		):
        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_		):
        lowercase_       :List[str]     =       tmp_path_factory.mktemp('''data'''		) / f"test.json.{extension}"
        lowercase_       :Dict     =       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[str]     =       f.read()
        with fsspec.open(UpperCamelCase_   ,							'''rb'''   ,							compression='''infer'''		) as f:
             lowercase_       :Dict     =       f.read()
        assert exported_content == original_content





from __future__ import annotations




def     UpperCamelCase    (      _a       ,		_a = None       ,		_a = None		)					->							None:
     '''simple docstring'''





     if start is None:
          lowercase_       :str     =       0

     if end is None:
          lowercase_       :str     =       len(_a		) - 1

     if start >= end:
          return

     lowercase_       :Dict     =       (start + end) // 2

     slowsort(_a       ,		_a       ,		_a		)
     slowsort(_a       ,		mid + 1       ,		_a		)

     if sequence[end] < sequence[mid]:
          lowercase_    ,      lowercase_       :List[Any]     =       sequence[mid], sequence[end]

     slowsort(_a       ,		_a       ,		end - 1		)


if __name__ == "__main__":
   from doctest import testmod

   testmod()

