





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 							SCREAMING_SNAKE_CASE__       (  snake_case__  :Tuple			,							snake_case__  :List[str]     )      ->						Union[str, Any]:
  assert isinstance(snake_case__			,							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 							SCREAMING_SNAKE_CASE__       (  snake_case__  :List[Any]			,							snake_case__  :Union[str, Any]			,							snake_case__  :Any     )      ->						int:
  _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(snake_case__			,							cache_dir=snake_case__			,							keep_in_memory=snake_case__     ).read()
  _check_json_dataset(snake_case__			,							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 							SCREAMING_SNAKE_CASE__       (  snake_case__  :Optional[Any]			,							snake_case__  :Optional[Any]			,							snake_case__  :Union[str, Any]     )      ->						Union[str, Any]:
  _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(snake_case__     ) for feature, dtype in features.items()}     ) if features is not None else None
  )
  _lowercase							= JsonDatasetReader(snake_case__			,							features=snake_case__			,							cache_dir=snake_case__     ).read()
  _check_json_dataset(snake_case__			,							snake_case__     )

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

@pytest.mark.parametrize('path_type'			,							[str, list]     )
def 							SCREAMING_SNAKE_CASE__       (  snake_case__  :Optional[Any]			,							snake_case__  :List[str]			,							snake_case__  :Tuple     )      ->						Dict:
  if issubclass(snake_case__			,							snake_case__     ):
    _lowercase							= jsonl_path
  elif issubclass(snake_case__			,							snake_case__     ):
    _lowercase							= [jsonl_path]
  _lowercase							= tmp_path / 'cache'
  _lowercase							= {'col_1': 'string', 'col_2': 'int64', 'col_3': 'float64'}
  _lowercase							= JsonDatasetReader(snake_case__			,							cache_dir=snake_case__     ).read()
  _check_json_dataset(snake_case__			,							snake_case__     )

def 							SCREAMING_SNAKE_CASE__       (  snake_case__  :List[str]			,							snake_case__  :List[Any]			,							snake_case__  :str=("train",)     )      ->						Union[str, Any]:
  assert isinstance(snake_case__			,							snake_case__     )
  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 							SCREAMING_SNAKE_CASE__       (  snake_case__  :Dict			,							snake_case__  :Optional[Any]			,							snake_case__  :Tuple     )      ->						Optional[int]:
  _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=snake_case__			,							keep_in_memory=snake_case__     ).read()
  _check_json_datasetdict(snake_case__			,							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 							SCREAMING_SNAKE_CASE__       (  snake_case__  :Any			,							snake_case__  :str			,							snake_case__  :List[str]     )      ->						int:
  _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(snake_case__     ) for feature, dtype in features.items()}     ) if features is not None else None
  )
  _lowercase							= JsonDatasetReader({'train': jsonl_path}			,							features=snake_case__			,							cache_dir=snake_case__     ).read()
  _check_json_datasetdict(snake_case__			,							snake_case__     )

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

def 							SCREAMING_SNAKE_CASE__       (  snake_case__  :List[str]     )      ->						Dict:
  return json.load(snake_case__     )

def 							SCREAMING_SNAKE_CASE__       (  snake_case__  :int     )      ->						List[str]:
  return [json.loads(snake_case__     ) for line in buffer]







class 				A_				:







       """simple docstring"""





       @pytest.mark.parametrize('lines, load_json_function'						,[(True, load_json_lines), (False, load_json)]  )
       def    __UpperCAmelCase (    self	:	Union[str, Any]						,__A	:	str						,__A	:	Any						,__A	:	Union[str, Any]  )						->					Tuple:
         with io.BytesIO() as buffer:
           JsonDatasetWriter(__A						,__A						,lines=__A  ).write()
           buffer.seek(0  )
           _lowercase							= load_json_function(__A  )
         assert isinstance(__A						,__A  )
         assert isinstance(exported_content[0]						,__A  )
         assert len(__A  ) == 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	:	int						,__A	:	List[str]						,__A	:	int						,__A	:	Optional[Any]						,__A	:	Dict						,__A	:	List[Any]  )						->					Dict:
         with io.BytesIO() as buffer:
           JsonDatasetWriter(__A						,__A						,lines=__A						,orient=__A  ).write()
           buffer.seek(0  )
           _lowercase							= load_json(__A  )
         assert isinstance(__A						,__A  )
         if keys:
           if container is dict:
             assert exported_content.keys() == keys
           else:
             assert exported_content[0].keys() == keys
         else:
           assert not hasattr(__A						,'keys'  ) and not hasattr(exported_content[0]						,'keys'  )
         if len_at:
           assert len(exported_content[len_at]  ) == 10
         else:
           assert len(__A  ) == 10







       @pytest.mark.parametrize('lines, load_json_function'						,[(True, load_json_lines), (False, load_json)]  )
       def    __UpperCAmelCase (    self	:	Dict						,__A	:	Optional[int]						,__A	:	Union[str, Any]						,__A	:	Dict  )						->					int:
         with io.BytesIO() as buffer:
           JsonDatasetWriter(__A						,__A						,lines=__A						,num_proc=2  ).write()
           buffer.seek(0  )
           _lowercase							= load_json_function(__A  )
         assert isinstance(__A						,__A  )
         assert isinstance(exported_content[0]						,__A  )
         assert len(__A  ) == 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	:	Optional[Any]						,__A	:	Dict						,__A	:	Union[str, Any]						,__A	:	Optional[Any]						,__A	:	str						,__A	:	Optional[int]  )						->					int:
         with io.BytesIO() as buffer:
           JsonDatasetWriter(__A						,__A						,lines=__A						,orient=__A						,num_proc=2  ).write()
           buffer.seek(0  )
           _lowercase							= load_json(__A  )
         assert isinstance(__A						,__A  )
         if keys:
           if container is dict:
             assert exported_content.keys() == keys
           else:
             assert exported_content[0].keys() == keys
         else:
           assert not hasattr(__A						,'keys'  ) and not hasattr(exported_content[0]						,'keys'  )
         if len_at:
           assert len(exported_content[len_at]  ) == 10
         else:
           assert len(__A  ) == 10







       def    __UpperCAmelCase (    self	:	str						,__A	:	Union[str, Any]  )						->					List[Any]:
         with pytest.raises(__A  ):
           with io.BytesIO() as buffer:
             JsonDatasetWriter(__A						,__A						,num_proc=0  )







       @pytest.mark.parametrize('compression, extension'						,[('gzip', 'gz'), ('bz2', 'bz2'), ('xz', 'xz')]  )
       def    __UpperCAmelCase (    self	:	List[str]						,__A	:	List[Any]						,__A	:	Any						,__A	:	Optional[int]						,__A	:	Dict						,__A	:	Optional[int]  )						->					Optional[int]:
         _lowercase							= tmp_path_factory.mktemp('data'  ) / F"""test.json.{extension}"""
         _lowercase							= str(shared_datadir / F"""test_file.json.{extension}"""  )
         JsonDatasetWriter(__A						,__A						,compression=__A  ).write()

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






import unittest

from transformers import is_torch_available
from transformers.testing_utils import require_torch


if is_torch_available():
      import torch

      from transformers.generation import DisjunctiveConstraint





@require_torch
class 				A_				(      unittest.TestCase	):







       """simple docstring"""





       def    __UpperCAmelCase (    self	:	Union[str, Any]  )						->					List[Any]:
         # For consistency across different places the DisjunctiveConstraint is called,
         # dc.token_ids is a list of integers. It is also initialized only by integers.

         _lowercase							= [[1, 2, 4], [1, 2, 3, 4]]
         _lowercase							= DisjunctiveConstraint(__A  )
         self.assertTrue(isinstance(dc.token_ids						,__A  )  )

         with self.assertRaises(__A  ):
           DisjunctiveConstraint(torch.LongTensor([[1, 2, 4], [1, 2, 3]]  )  )

         with self.assertRaises(__A  ):
           DisjunctiveConstraint([torch.LongTensor([1, 2, 4]  ), torch.LongTensor([1, 2, 3, 4, 5]  )]  )







       def    __UpperCAmelCase (    self	:	List[Any]  )						->					str:
         # We can't have constraints that are complete subsets of another. This leads to a preverse
         # interpretation of "constraint fulfillment": does generating [1,2,3] fulfill the constraint?
         # It would mean that it generated [1,2] which fulfills it, but it's in the middle of potentially
         # fulfilling [1,2,3,4]. If we believe that [1,2,3] does fulfill the constraint, then the algorithm
         # will necessarily never reach [1,2,3,4], giving users a false sense of control (better to just not allow it).
         _lowercase							= [[1, 2], [1, 2, 3, 4]]

         with self.assertRaises(__A  ):
           DisjunctiveConstraint(__A  )  # fails here







       def    __UpperCAmelCase (    self	:	Dict  )						->					Optional[Any]:
         _lowercase							= [[1, 2, 3], [1, 2, 4]]

         _lowercase							= DisjunctiveConstraint(__A  )

         _lowercase						,		_lowercase						,		_lowercase							= dc.update(1  )
         _lowercase							= stepped is True and completed is False and reset is False
         self.assertTrue(__A  )
         self.assertTrue(not dc.completed  )
         self.assertTrue(dc.current_seq == [1]  )

         _lowercase						,		_lowercase						,		_lowercase							= dc.update(2  )
         _lowercase							= stepped is True and completed is False and reset is False
         self.assertTrue(__A  )
         self.assertTrue(not dc.completed  )
         self.assertTrue(dc.current_seq == [1, 2]  )

         _lowercase						,		_lowercase						,		_lowercase							= dc.update(3  )
         _lowercase							= stepped is True and completed is True and reset is False
         self.assertTrue(__A  )
         self.assertTrue(dc.completed  )  # Completed!
         self.assertTrue(dc.current_seq == [1, 2, 3]  )







       def    __UpperCAmelCase (    self	:	List[Any]  )						->					Any:
         _lowercase							= [[1, 2, 3], [1, 2, 4, 5], [1, 2, 5]]

         _lowercase							= DisjunctiveConstraint(__A  )

         _lowercase						,		_lowercase						,		_lowercase							= dc.update(1  )
         self.assertTrue(not dc.completed  )
         self.assertTrue(dc.current_seq == [1]  )

         _lowercase						,		_lowercase						,		_lowercase							= dc.update(2  )
         self.assertTrue(not dc.completed  )
         self.assertTrue(dc.current_seq == [1, 2]  )

         _lowercase						,		_lowercase						,		_lowercase							= dc.update(4  )
         self.assertTrue(not dc.completed  )
         self.assertTrue(dc.current_seq == [1, 2, 4]  )

         _lowercase						,		_lowercase						,		_lowercase							= dc.update(5  )
         self.assertTrue(dc.completed  )  # Completed!
         self.assertTrue(dc.current_seq == [1, 2, 4, 5]  )

         dc.reset()

         _lowercase						,		_lowercase						,		_lowercase							= dc.update(1  )
         self.assertTrue(not dc.completed  )
         self.assertTrue(dc.remaining() == 3  )
         self.assertTrue(dc.current_seq == [1]  )

         _lowercase						,		_lowercase						,		_lowercase							= dc.update(2  )
         self.assertTrue(not dc.completed  )
         self.assertTrue(dc.remaining() == 2  )
         self.assertTrue(dc.current_seq == [1, 2]  )

         _lowercase						,		_lowercase						,		_lowercase							= dc.update(5  )
         self.assertTrue(dc.completed  )  # Completed!
         self.assertTrue(dc.remaining() == 0  )
         self.assertTrue(dc.current_seq == [1, 2, 5]  )