"""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 					__a   (		A    ,				A      ):






     '''simple docstring'''



     assert isinstance(lowerCAmelCase__    ,				lowerCAmelCase__      )
     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 					__a   (		A    ,				A    ,				A      ):






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







@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 					__a   (		A    ,				A    ,				A      ):






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







@pytest.mark.parametrize(
    "features"    ,				[
        None,
        {"col_3": "float64", "col_1": "string", "col_2": "int64"},
    ]    ,				)
def 					__a   (		A    ,				A    ,				A      ):






     '''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(lowerCAmelCase__      ) for feature, dtype in features.items()}      ) if features is not None else None
     )
     lowercase__				 =							JsonDatasetReader(lowerCAmelCase__    ,				features=lowerCAmelCase__    ,				cache_dir=lowerCAmelCase__      ).read()
     assert isinstance(lowerCAmelCase__    ,				lowerCAmelCase__      )
     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 					__a   (		A    ,				A      ):






     '''simple docstring'''



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






     '''simple docstring'''



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







@pytest.mark.parametrize("path_type"    ,				[str, list]      )
def 					__a   (		A    ,				A    ,				A      ):






     '''simple docstring'''



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







def 					__a   (		A    ,				A    ,				A=("train",)      ):






     '''simple docstring'''



     assert isinstance(lowerCAmelCase__    ,				lowerCAmelCase__      )
     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 					__a   (		A    ,				A    ,				A      ):






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







@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 					__a   (		A    ,				A    ,				A      ):






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







@pytest.mark.parametrize("split"    ,				[None, NamedSplit("train"      ), "train", "test"]      )
def 					__a   (		A    ,				A    ,				A      ):






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







def 					__a   (		A      ):






     '''simple docstring'''



     return json.load(lowerCAmelCase__      )







def 					__a   (		A      ):






     '''simple docstring'''



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






class    a__    :


   @pytest.mark.parametrize("lines, load_json_function",				[(True, load_json_lines), (False, load_json)]  )
   def 	snake_case__	(     self,				_UpperCAmelCase,				_UpperCAmelCase,				_UpperCAmelCase  ):
        '''simple docstring'''



        with io.BytesIO() as buffer:
             JsonDatasetWriter(UpperCamelCase__,				UpperCamelCase__,				lines=UpperCamelCase__  ).write()
             buffer.seek(0  )
             lowercase__				 =							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 	snake_case__	(     self,				_UpperCAmelCase,				_UpperCAmelCase,				_UpperCAmelCase,				_UpperCAmelCase,				_UpperCAmelCase  ):
        '''simple docstring'''



        with io.BytesIO() as buffer:
             JsonDatasetWriter(UpperCamelCase__,				UpperCamelCase__,				lines=UpperCamelCase__,				orient=UpperCamelCase__  ).write()
             buffer.seek(0  )
             lowercase__				 =							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 	snake_case__	(     self,				_UpperCAmelCase,				_UpperCAmelCase,				_UpperCAmelCase  ):
        '''simple docstring'''



        with io.BytesIO() as buffer:
             JsonDatasetWriter(UpperCamelCase__,				UpperCamelCase__,				lines=UpperCamelCase__,				num_proc=2  ).write()
             buffer.seek(0  )
             lowercase__				 =							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 	snake_case__	(     self,				_UpperCAmelCase,				_UpperCAmelCase,				_UpperCAmelCase,				_UpperCAmelCase,				_UpperCAmelCase  ):
        '''simple docstring'''



        with io.BytesIO() as buffer:
             JsonDatasetWriter(UpperCamelCase__,				UpperCamelCase__,				lines=UpperCamelCase__,				orient=UpperCamelCase__,				num_proc=2  ).write()
             buffer.seek(0  )
             lowercase__				 =							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 	snake_case__	(     self,				_UpperCAmelCase  ):
        '''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 	snake_case__	(     self,				_UpperCAmelCase,				_UpperCAmelCase,				_UpperCAmelCase,				_UpperCAmelCase,				_UpperCAmelCase  ):
        '''simple docstring'''



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

        with fsspec.open(UpperCamelCase__,				"rb",				compression="infer"  ) as f:
             lowercase__				 =							f.read()
        with fsspec.open(UpperCamelCase__,				"rb",				compression="infer"  ) as f:
             lowercase__				 =							f.read()
        assert exported_content == original_content





"""simple docstring"""


lowerCAmelCase_:					Any     =					"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/"







def 					__a   (		A      ):






     '''simple docstring'''



     if not isinstance(A    ,				A      ):
          lowercase__				 =							f'''a bytes-like object is required, not \'{data.__class__.__name__}\''''
          raise TypeError(A      )

     lowercase__				 =							"".join(bin(A      )[2:].zfill(8      ) for byte in data      )

     lowercase__				 =							len(A      ) % 6 != 0

     if padding_needed:
          # The padding that will be added later
          lowercase__				 =							b"=" * ((6 - len(A      ) % 6) // 2)

          # Append binary_stream with arbitrary binary digits (0's by default) to make its
          # length a multiple of 6.
          binary_stream += "0" * (6 - len(A      ) % 6)
     else:
          lowercase__				 =							b""

     # Encode every 6 binary digits to their corresponding Base64 character
     return (
         "".join(
             B64_CHARSET[int(binary_stream[index : index + 6]    ,				2      )]
             for index in range(0    ,				len(A      )    ,				6      )      ).encode()
         + padding
     )







def 					__a   (		A      ):






     '''simple docstring'''



     if not isinstance(A    ,				A      ) and not isinstance(A    ,				A      ):
          lowercase__				 =							(
              "argument should be a bytes-like object or ASCII string, "
              f'''not \'{encoded_data.__class__.__name__}\''''
          )
          raise TypeError(A      )

     # In case encoded_data is a bytes-like object, make sure it contains only
     # ASCII characters so we convert it to a string object
     if isinstance(A    ,				A      ):
          try:
               lowercase__				 =							encoded_data.decode("utf-8"      )
          except UnicodeDecodeError:
               raise ValueError("base64 encoded data should only contain ASCII characters"      )

     lowercase__				 =							encoded_data.count("="      )

     # Check if the encoded string contains non base64 characters
     if padding:
          assert all(
              char in B64_CHARSET for char in encoded_data[:-padding]      ), "Invalid base64 character(s) found."
     else:
          assert all(
              char in B64_CHARSET for char in encoded_data      ), "Invalid base64 character(s) found."

     # Check the padding
     assert len(A      ) % 4 == 0 and padding < 3, "Incorrect padding"

     if padding:
          # Remove padding if there is one
          lowercase__				 =							encoded_data[:-padding]

          lowercase__				 =							"".join(
              bin(B64_CHARSET.index(A      )      )[2:].zfill(6      ) for char in encoded_data      )[: -padding * 2]
     else:
          lowercase__				 =							"".join(
              bin(B64_CHARSET.index(A      )      )[2:].zfill(6      ) for char in encoded_data      )

     lowercase__				 =							[
         int(binary_stream[index : index + 8]    ,				2      )
         for index in range(0    ,				len(A      )    ,				8      )
     ]

     return bytes(A      )


if __name__ == "__main__":
 import doctest

 doctest.testmod()




