
"""simple docstring"""


import gc
import random
import tempfile
import unittest

import numpy as np
import torch
from PIL import Image
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMInverseScheduler,
    DDIMScheduler,
    DPMSolverMultistepInverseScheduler,
    DPMSolverMultistepScheduler,
    StableDiffusionDiffEditPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import load_image, slow
from diffusers.utils.testing_utils import enable_full_determinism, floats_tensor, require_torch_gpu, torch_device

from ..pipeline_params import TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS, TEXT_GUIDED_IMAGE_INPAINTING_PARAMS
from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()




class  __lowerCamelCase  ( lowerCAmelCase						,	lowerCAmelCase						,	unittest.TestCase			):
				a__:      Optional[int]          =    StableDiffusionDiffEditPipeline
				a__:      Any          =    TEXT_GUIDED_IMAGE_INPAINTING_PARAMS - {'height', 'width', 'image'} | {'image_latents'}
				a__:      Optional[int]          =    TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS - {'image'} | {'image_latents'}
				a__:      str          =    frozenset(
				    []			)  # TO-DO: update image_params once pipeline is refactored with VaeImageProcessor.preprocess
				a__:      Dict          =    frozenset([]			)




				def        UpperCAmelCase__						(				self   ):
							torch.manual_seed(0   )
							lowerCamelCase_					=  UNetaDConditionModel(
							    block_out_channels=(32, 64)							,						layers_per_block=2							,						sample_size=32							,						in_channels=4							,						out_channels=4							,						down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D''')							,						up_block_types=('''CrossAttnUpBlock2D''', '''UpBlock2D''')							,						cross_attention_dim=32							,						attention_head_dim=(2, 4)							,						use_linear_projection=UpperCAmelCase							,						)
							lowerCamelCase_					=  DDIMScheduler(
							    beta_start=0.0_0_0_8_5							,						beta_end=0.0_1_2							,						beta_schedule='''scaled_linear'''							,						clip_sample=UpperCAmelCase							,						set_alpha_to_one=UpperCAmelCase							,						)
							lowerCamelCase_					=  DDIMInverseScheduler(
							    beta_start=0.0_0_0_8_5							,						beta_end=0.0_1_2							,						beta_schedule='''scaled_linear'''							,						clip_sample=UpperCAmelCase							,						set_alpha_to_zero=UpperCAmelCase							,						)
							torch.manual_seed(0   )
							lowerCamelCase_					=  AutoencoderKL(
							    block_out_channels=[32, 64]							,						in_channels=3							,						out_channels=3							,						down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D''']							,						up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D''']							,						latent_channels=4							,						sample_size=128							,						)
							torch.manual_seed(0   )
							lowerCamelCase_					=  CLIPTextConfig(
							    bos_token_id=0							,						eos_token_id=2							,						hidden_size=32							,						intermediate_size=37							,						layer_norm_eps=1e-0_5							,						num_attention_heads=4							,						num_hidden_layers=5							,						pad_token_id=1							,						vocab_size=1000							,						hidden_act='''gelu'''							,						projection_dim=512							,						)
							lowerCamelCase_					=  CLIPTextModel(UpperCAmelCase   )
							lowerCamelCase_					=  CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''   )

							lowerCamelCase_					=  {
							    '''unet''': unet,
							    '''scheduler''': scheduler,
							    '''inverse_scheduler''': inverse_scheduler,
							    '''vae''': vae,
							    '''text_encoder''': text_encoder,
							    '''tokenizer''': tokenizer,
							    '''safety_checker''': None,
							    '''feature_extractor''': None,
							}

							return components




				def        UpperCAmelCase__						(				self							,						UpperCAmelCase							,						UpperCAmelCase=0   ):
							lowerCamelCase_					=  floats_tensor((1, 16, 16)							,						rng=random.Random(UpperCAmelCase   )   ).to(UpperCAmelCase   )
							lowerCamelCase_					=  floats_tensor((1, 2, 4, 16, 16)							,						rng=random.Random(UpperCAmelCase   )   ).to(UpperCAmelCase   )
							if str(UpperCAmelCase   ).startswith('''mps'''   ):
										lowerCamelCase_					=  torch.manual_seed(UpperCAmelCase   )
							else:
										lowerCamelCase_					=  torch.Generator(device=UpperCAmelCase   ).manual_seed(UpperCAmelCase   )
							lowerCamelCase_					=  {
							    '''prompt''': '''a dog and a newt''',
							    '''mask_image''': mask,
							    '''image_latents''': latents,
							    '''generator''': generator,
							    '''num_inference_steps''': 2,
							    '''inpaint_strength''': 1.0,
							    '''guidance_scale''': 6.0,
							    '''output_type''': '''numpy''',
							}

							return inputs




				def        UpperCAmelCase__						(				self							,						UpperCAmelCase							,						UpperCAmelCase=0   ):
							lowerCamelCase_					=  floats_tensor((1, 3, 32, 32)							,						rng=random.Random(UpperCAmelCase   )   ).to(UpperCAmelCase   )
							lowerCamelCase_					=  image.cpu().permute(0							,						2							,						3							,						1   )[0]
							lowerCamelCase_					=  Image.fromarray(np.uinta(UpperCAmelCase   )   ).convert('''RGB'''   )
							if str(UpperCAmelCase   ).startswith('''mps'''   ):
										lowerCamelCase_					=  torch.manual_seed(UpperCAmelCase   )
							else:
										lowerCamelCase_					=  torch.Generator(device=UpperCAmelCase   ).manual_seed(UpperCAmelCase   )
							lowerCamelCase_					=  {
							    '''image''': image,
							    '''source_prompt''': '''a cat and a frog''',
							    '''target_prompt''': '''a dog and a newt''',
							    '''generator''': generator,
							    '''num_inference_steps''': 2,
							    '''num_maps_per_mask''': 2,
							    '''mask_encode_strength''': 1.0,
							    '''guidance_scale''': 6.0,
							    '''output_type''': '''numpy''',
							}

							return inputs




				def        UpperCAmelCase__						(				self							,						UpperCAmelCase							,						UpperCAmelCase=0   ):
							lowerCamelCase_					=  floats_tensor((1, 3, 32, 32)							,						rng=random.Random(UpperCAmelCase   )   ).to(UpperCAmelCase   )
							lowerCamelCase_					=  image.cpu().permute(0							,						2							,						3							,						1   )[0]
							lowerCamelCase_					=  Image.fromarray(np.uinta(UpperCAmelCase   )   ).convert('''RGB'''   )
							if str(UpperCAmelCase   ).startswith('''mps'''   ):
										lowerCamelCase_					=  torch.manual_seed(UpperCAmelCase   )
							else:
										lowerCamelCase_					=  torch.Generator(device=UpperCAmelCase   ).manual_seed(UpperCAmelCase   )
							lowerCamelCase_					=  {
							    '''image''': image,
							    '''prompt''': '''a cat and a frog''',
							    '''generator''': generator,
							    '''num_inference_steps''': 2,
							    '''inpaint_strength''': 1.0,
							    '''guidance_scale''': 6.0,
							    '''decode_latents''': True,
							    '''output_type''': '''numpy''',
							}
							return inputs




				def        UpperCAmelCase__						(				self   ):
							if not hasattr(self.pipeline_class							,						'''_optional_components'''   ):
										return

							lowerCamelCase_					=  self.get_dummy_components()
							lowerCamelCase_					=  self.pipeline_class(**UpperCAmelCase   )
							pipe.to(UpperCAmelCase   )
							pipe.set_progress_bar_config(disable=UpperCAmelCase   )

							# set all optional components to None and update pipeline config accordingly
							for optional_component in pipe._optional_components:
										setattr(UpperCAmelCase							,						UpperCAmelCase							,						UpperCAmelCase   )
							pipe.register_modules(**{optional_component: None for optional_component in pipe._optional_components}   )

							lowerCamelCase_					=  self.get_dummy_inputs(UpperCAmelCase   )
							lowerCamelCase_					=  pipe(**UpperCAmelCase   )[0]

							with tempfile.TemporaryDirectory() as tmpdir:
										pipe.save_pretrained(UpperCAmelCase   )
										lowerCamelCase_					=  self.pipeline_class.from_pretrained(UpperCAmelCase   )
										pipe_loaded.to(UpperCAmelCase   )
										pipe_loaded.set_progress_bar_config(disable=UpperCAmelCase   )

							for optional_component in pipe._optional_components:
										self.assertTrue(
										    getattr(UpperCAmelCase							,						UpperCAmelCase   ) is None							,						f"`{optional_component}` did not stay set to None after loading."							,						)

							lowerCamelCase_					=  self.get_dummy_inputs(UpperCAmelCase   )
							lowerCamelCase_					=  pipe_loaded(**UpperCAmelCase   )[0]

							lowerCamelCase_					=  np.abs(output - output_loaded   ).max()
							self.assertLess(UpperCAmelCase							,						1e-4   )




				def        UpperCAmelCase__						(				self   ):
							lowerCamelCase_					=  '''cpu'''

							lowerCamelCase_					=  self.get_dummy_components()
							lowerCamelCase_					=  self.pipeline_class(**UpperCAmelCase   )
							pipe.to(UpperCAmelCase   )
							pipe.set_progress_bar_config(disable=UpperCAmelCase   )

							lowerCamelCase_					=  self.get_dummy_mask_inputs(UpperCAmelCase   )
							lowerCamelCase_					=  pipe.generate_mask(**UpperCAmelCase   )
							lowerCamelCase_					=  mask[0, -3:, -3:]

							self.assertEqual(mask.shape							,						(1, 16, 16)   )
							lowerCamelCase_					=  np.array([0] * 9   )
							lowerCamelCase_					=  np.abs(mask_slice.flatten() - expected_slice   ).max()
							self.assertLessEqual(UpperCAmelCase							,						1e-3   )
							self.assertEqual(mask[0, -3, -4]							,						0   )




				def        UpperCAmelCase__						(				self   ):
							lowerCamelCase_					=  '''cpu'''

							lowerCamelCase_					=  self.get_dummy_components()
							lowerCamelCase_					=  self.pipeline_class(**UpperCAmelCase   )
							pipe.to(UpperCAmelCase   )
							pipe.set_progress_bar_config(disable=UpperCAmelCase   )

							lowerCamelCase_					=  self.get_dummy_inversion_inputs(UpperCAmelCase   )
							lowerCamelCase_					=  pipe.invert(**UpperCAmelCase   ).images
							lowerCamelCase_					=  image[0, -1, -3:, -3:]

							self.assertEqual(image.shape							,						(2, 32, 32, 3)   )
							lowerCamelCase_					=  np.array(
							    [0.5_1_5_0, 0.5_1_3_4, 0.5_0_4_3, 0.5_3_7_6, 0.4_6_9_4, 0.5_1_0_5_0, 0.5_0_1_5, 0.4_4_0_7, 0.4_7_9_9]							,						)
							lowerCamelCase_					=  np.abs(image_slice.flatten() - expected_slice   ).max()
							self.assertLessEqual(UpperCAmelCase							,						1e-3   )




				def        UpperCAmelCase__						(				self   ):
							super().test_inference_batch_single_identical(expected_max_diff=5e-3   )





				def        UpperCAmelCase__						(				self   ):
							lowerCamelCase_					=  '''cpu'''

							lowerCamelCase_					=  self.get_dummy_components()

							lowerCamelCase_					=  {'''beta_start''': 0.0_0_0_8_5, '''beta_end''': 0.0_1_2, '''beta_schedule''': '''scaled_linear'''}
							lowerCamelCase_					=  DPMSolverMultistepScheduler(**UpperCAmelCase   )
							lowerCamelCase_					=  DPMSolverMultistepInverseScheduler(**UpperCAmelCase   )

							lowerCamelCase_					=  self.pipeline_class(**UpperCAmelCase   )
							pipe.to(UpperCAmelCase   )
							pipe.set_progress_bar_config(disable=UpperCAmelCase   )

							lowerCamelCase_					=  self.get_dummy_inversion_inputs(UpperCAmelCase   )
							lowerCamelCase_					=  pipe.invert(**UpperCAmelCase   ).images
							lowerCamelCase_					=  image[0, -1, -3:, -3:]

							self.assertEqual(image.shape							,						(2, 32, 32, 3)   )
							lowerCamelCase_					=  np.array(
							    [0.5_1_5_0, 0.5_1_3_4, 0.5_0_4_3, 0.5_3_7_6, 0.4_6_9_4, 0.5_1_0_5_0, 0.5_0_1_5, 0.4_4_0_7, 0.4_7_9_9]							,						)
							lowerCamelCase_					=  np.abs(image_slice.flatten() - expected_slice   ).max()
							self.assertLessEqual(UpperCAmelCase							,						1e-3   )




@require_torch_gpu
@slow
class  __lowerCamelCase  ( unittest.TestCase			):




				def        UpperCAmelCase__						(				self   ):
							super().tearDown()
							gc.collect()
							torch.cuda.empty_cache()




				@classmethod
				def        UpperCAmelCase__						(				cls   ):
							lowerCamelCase_					=  load_image(
							    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/diffedit/fruit.png'''   )

							lowerCamelCase_					=  raw_image.convert('''RGB'''   ).resize((768, 768)   )

							lowerCamelCase_					=  raw_image




				def        UpperCAmelCase__						(				self   ):
							lowerCamelCase_					=  torch.manual_seed(0   )

							lowerCamelCase_					=  StableDiffusionDiffEditPipeline.from_pretrained(
							    '''stabilityai/stable-diffusion-2-1'''							,						safety_checker=UpperCAmelCase							,						torch_dtype=torch.floataa   )
							lowerCamelCase_					=  DDIMScheduler.from_config(pipe.scheduler.config   )
							lowerCamelCase_					=  DDIMInverseScheduler.from_config(pipe.scheduler.config   )
							pipe.enable_model_cpu_offload()
							pipe.set_progress_bar_config(disable=UpperCAmelCase   )

							lowerCamelCase_					=  '''a bowl of fruit'''
							lowerCamelCase_					=  '''a bowl of pears'''

							lowerCamelCase_					=  pipe.generate_mask(
							    image=self.raw_image							,						source_prompt=UpperCAmelCase							,						target_prompt=UpperCAmelCase							,						generator=UpperCAmelCase							,						)

							lowerCamelCase_					=  pipe.invert(
							    prompt=UpperCAmelCase							,						image=self.raw_image							,						inpaint_strength=0.7							,						generator=UpperCAmelCase   ).latents

							lowerCamelCase_					=  pipe(
							    prompt=UpperCAmelCase							,						mask_image=UpperCAmelCase							,						image_latents=UpperCAmelCase							,						generator=UpperCAmelCase							,						negative_prompt=UpperCAmelCase							,						inpaint_strength=0.7							,						output_type='''numpy'''							,						).images[0]

							lowerCamelCase_					=  (
							    np.array(
							        load_image(
							            '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
							            '''/diffedit/pears.png'''   ).resize((768, 768)   )   )
							    / 255
							)
							assert np.abs((expected_image - image).max()   ) < 5e-1





				def        UpperCAmelCase__						(				self   ):
							lowerCamelCase_					=  torch.manual_seed(0   )

							lowerCamelCase_					=  StableDiffusionDiffEditPipeline.from_pretrained(
							    '''stabilityai/stable-diffusion-2-1'''							,						safety_checker=UpperCAmelCase							,						torch_dtype=torch.floataa   )
							lowerCamelCase_					=  DPMSolverMultistepScheduler.from_config(pipe.scheduler.config   )
							lowerCamelCase_					=  DPMSolverMultistepInverseScheduler.from_config(pipe.scheduler.config   )
							pipe.enable_model_cpu_offload()
							pipe.set_progress_bar_config(disable=UpperCAmelCase   )

							lowerCamelCase_					=  '''a bowl of fruit'''
							lowerCamelCase_					=  '''a bowl of pears'''

							lowerCamelCase_					=  pipe.generate_mask(
							    image=self.raw_image							,						source_prompt=UpperCAmelCase							,						target_prompt=UpperCAmelCase							,						generator=UpperCAmelCase							,						)

							lowerCamelCase_					=  pipe.invert(
							    prompt=UpperCAmelCase							,						image=self.raw_image							,						inpaint_strength=0.7							,						generator=UpperCAmelCase							,						num_inference_steps=25							,						).latents

							lowerCamelCase_					=  pipe(
							    prompt=UpperCAmelCase							,						mask_image=UpperCAmelCase							,						image_latents=UpperCAmelCase							,						generator=UpperCAmelCase							,						negative_prompt=UpperCAmelCase							,						inpaint_strength=0.7							,						num_inference_steps=25							,						output_type='''numpy'''							,						).images[0]

							lowerCamelCase_					=  (
							    np.array(
							        load_image(
							            '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
							            '''/diffedit/pears.png'''   ).resize((768, 768)   )   )
							    / 255
							)
							assert np.abs((expected_image - image).max()   ) < 5e-1










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  (	UpperCamelCase__   :   Any       ,			UpperCamelCase__   :   Optional[Any]					)      ->    str:

						"""simple docstring"""





						assert isinstance(UpperCamelCase__       ,			UpperCamelCase__					)
						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  (	UpperCamelCase__   :   Optional[int]       ,			UpperCamelCase__   :   Optional[int]       ,			UpperCamelCase__   :   List[Any]					)      ->    Tuple:

						"""simple docstring"""





						__SCREAMING_SNAKE_CASE:				Dict           =   tmp_path / '''cache'''
						__SCREAMING_SNAKE_CASE:				Dict           =   {'''col_1''': '''string''', '''col_2''': '''int64''', '''col_3''': '''float64'''}
						with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
												__SCREAMING_SNAKE_CASE:				Tuple           =   JsonDatasetReader(UpperCamelCase__       ,			cache_dir=UpperCamelCase__       ,			keep_in_memory=UpperCamelCase__					).read()
						_check_json_dataset(UpperCamelCase__       ,			UpperCamelCase__					)






@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  (	UpperCamelCase__   :   List[str]       ,			UpperCamelCase__   :   int       ,			UpperCamelCase__   :   Optional[int]					)      ->    int:

						"""simple docstring"""





						__SCREAMING_SNAKE_CASE:				List[str]           =   tmp_path / '''cache'''
						__SCREAMING_SNAKE_CASE:				Tuple           =   {'''col_1''': '''string''', '''col_2''': '''int64''', '''col_3''': '''float64'''}
						__SCREAMING_SNAKE_CASE:				Any           =   features.copy() if features else default_expected_features
						__SCREAMING_SNAKE_CASE:				List[Any]           =   (
						    Features({feature: Value(UpperCamelCase__					) for feature, dtype in features.items()}					) if features is not None else None
						)
						__SCREAMING_SNAKE_CASE:				Tuple           =   JsonDatasetReader(UpperCamelCase__       ,			features=UpperCamelCase__       ,			cache_dir=UpperCamelCase__					).read()
						_check_json_dataset(UpperCamelCase__       ,			UpperCamelCase__					)






@pytest.mark.parametrize(
    '''features'''       ,			[
        None,
        {'''col_3''': '''float64''', '''col_1''': '''string''', '''col_2''': '''int64'''},
    ]       ,			)
def 		lowerCAmelCase  (	UpperCamelCase__   :   Union[str, Any]       ,			UpperCamelCase__   :   List[str]       ,			UpperCamelCase__   :   Optional[int]					)      ->    Dict:

						"""simple docstring"""





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

						"""simple docstring"""





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

						"""simple docstring"""





						__SCREAMING_SNAKE_CASE:				List[Any]           =   tmp_path / '''cache'''
						__SCREAMING_SNAKE_CASE:				Optional[Any]           =   {'''col_1''': '''string''', '''col_2''': '''int64''', '''col_3''': '''float64'''}
						__SCREAMING_SNAKE_CASE:				Dict           =   JsonDatasetReader(UpperCamelCase__       ,			cache_dir=UpperCamelCase__       ,			split=UpperCamelCase__					).read()
						_check_json_dataset(UpperCamelCase__       ,			UpperCamelCase__					)
						assert dataset.split == split if split else "train"






@pytest.mark.parametrize('''path_type'''       ,			[str, list]					)
def 		lowerCAmelCase  (	UpperCamelCase__   :   Optional[Any]       ,			UpperCamelCase__   :   Union[str, Any]       ,			UpperCamelCase__   :   str					)      ->    Any:

						"""simple docstring"""





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






def 		lowerCAmelCase  (	UpperCamelCase__   :   Optional[int]       ,			UpperCamelCase__   :   int       ,			UpperCamelCase__   :   Union[str, Any]=("train",)					)      ->    List[str]:

						"""simple docstring"""





						assert isinstance(UpperCamelCase__       ,			UpperCamelCase__					)
						for split in splits:
												__SCREAMING_SNAKE_CASE:				Tuple           =   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  (	UpperCamelCase__   :   List[Any]       ,			UpperCamelCase__   :   Union[str, Any]       ,			UpperCamelCase__   :   Tuple					)      ->    Dict:

						"""simple docstring"""





						__SCREAMING_SNAKE_CASE:				Optional[Any]           =   tmp_path / '''cache'''
						__SCREAMING_SNAKE_CASE:				List[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():
												__SCREAMING_SNAKE_CASE:				Dict           =   JsonDatasetReader({'''train''': jsonl_path}       ,			cache_dir=UpperCamelCase__       ,			keep_in_memory=UpperCamelCase__					).read()
						_check_json_datasetdict(UpperCamelCase__       ,			UpperCamelCase__					)






@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  (	UpperCamelCase__   :   Dict       ,			UpperCamelCase__   :   Tuple       ,			UpperCamelCase__   :   List[str]					)      ->    Union[str, Any]:

						"""simple docstring"""





						__SCREAMING_SNAKE_CASE:				Optional[Any]           =   tmp_path / '''cache'''
						__SCREAMING_SNAKE_CASE:				Union[str, Any]           =   {'''col_1''': '''string''', '''col_2''': '''int64''', '''col_3''': '''float64'''}
						__SCREAMING_SNAKE_CASE:				str           =   features.copy() if features else default_expected_features
						__SCREAMING_SNAKE_CASE:				Optional[int]           =   (
						    Features({feature: Value(UpperCamelCase__					) for feature, dtype in features.items()}					) if features is not None else None
						)
						__SCREAMING_SNAKE_CASE:				Dict           =   JsonDatasetReader({'''train''': jsonl_path}       ,			features=UpperCamelCase__       ,			cache_dir=UpperCamelCase__					).read()
						_check_json_datasetdict(UpperCamelCase__       ,			UpperCamelCase__					)






@pytest.mark.parametrize('''split'''       ,			[None, NamedSplit('''train'''					), '''train''', '''test''']					)
def 		lowerCAmelCase  (	UpperCamelCase__   :   List[Any]       ,			UpperCamelCase__   :   int       ,			UpperCamelCase__   :   int					)      ->    Optional[Any]:

						"""simple docstring"""





						if split:
												__SCREAMING_SNAKE_CASE:				str           =   {split: jsonl_path}
						else:
												__SCREAMING_SNAKE_CASE:				Dict           =   '''train'''
												__SCREAMING_SNAKE_CASE:				Tuple           =   {'''train''': jsonl_path, '''test''': jsonl_path}
						__SCREAMING_SNAKE_CASE:				Tuple           =   tmp_path / '''cache'''
						__SCREAMING_SNAKE_CASE:				str           =   {'''col_1''': '''string''', '''col_2''': '''int64''', '''col_3''': '''float64'''}
						__SCREAMING_SNAKE_CASE:				int           =   JsonDatasetReader(UpperCamelCase__       ,			cache_dir=UpperCamelCase__					).read()
						_check_json_datasetdict(UpperCamelCase__       ,			UpperCamelCase__       ,			splits=list(path.keys()					)					)
						assert all(dataset[split].split == split for split in path.keys()					)






def 		lowerCAmelCase  (	UpperCamelCase__   :   Dict					)      ->    Optional[int]:

						"""simple docstring"""





						return json.load(UpperCamelCase__					)






def 		lowerCAmelCase  (	UpperCamelCase__   :   Any					)      ->    Any:

						"""simple docstring"""





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





class 				a			:






				@pytest.mark.parametrize('''lines, load_json_function'''			, [(True, load_json_lines), (False, load_json)]					)
				def 			snake_case_    (  self			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase					):

										"""simple docstring"""




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

										"""simple docstring"""




										with io.BytesIO() as buffer:
																JsonDatasetWriter(_lowerCAmelCase			, _lowerCAmelCase			, lines=_lowerCAmelCase			, orient=_lowerCAmelCase					).write()
																buffer.seek(0					)
																__SCREAMING_SNAKE_CASE:				Optional[int]           =   load_json(_lowerCAmelCase					)
										assert isinstance(_lowerCAmelCase			, _lowerCAmelCase					)
										if keys:
																if container is dict:
																						assert exported_content.keys() == keys
																else:
																						assert exported_content[0].keys() == keys
										else:
																assert not hasattr(_lowerCAmelCase			, '''keys'''					) and not hasattr(exported_content[0]			, '''keys'''					)
										if len_at:
																assert len(exported_content[len_at]					) == 10
										else:
																assert len(_lowerCAmelCase					) == 10






				@pytest.mark.parametrize('''lines, load_json_function'''			, [(True, load_json_lines), (False, load_json)]					)
				def 			snake_case_    (  self			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase					):

										"""simple docstring"""




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

										"""simple docstring"""




										with io.BytesIO() as buffer:
																JsonDatasetWriter(_lowerCAmelCase			, _lowerCAmelCase			, lines=_lowerCAmelCase			, orient=_lowerCAmelCase			, num_proc=2					).write()
																buffer.seek(0					)
																__SCREAMING_SNAKE_CASE:				Dict           =   load_json(_lowerCAmelCase					)
										assert isinstance(_lowerCAmelCase			, _lowerCAmelCase					)
										if keys:
																if container is dict:
																						assert exported_content.keys() == keys
																else:
																						assert exported_content[0].keys() == keys
										else:
																assert not hasattr(_lowerCAmelCase			, '''keys'''					) and not hasattr(exported_content[0]			, '''keys'''					)
										if len_at:
																assert len(exported_content[len_at]					) == 10
										else:
																assert len(_lowerCAmelCase					) == 10






				def 			snake_case_    (  self			, _lowerCAmelCase					):

										"""simple docstring"""




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






				@pytest.mark.parametrize('''compression, extension'''			, [('''gzip''', '''gz'''), ('''bz2''', '''bz2'''), ('''xz''', '''xz''')]					)
				def 			snake_case_    (  self			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase					):

										"""simple docstring"""




										__SCREAMING_SNAKE_CASE:				Any           =   tmp_path_factory.mktemp('''data'''					) / f"""test.json.{extension}"""
										__SCREAMING_SNAKE_CASE:				Any           =   str(shared_datadir / f"""test_file.json.{extension}"""					)
										JsonDatasetWriter(_lowerCAmelCase			, _lowerCAmelCase			, compression=_lowerCAmelCase					).write()

										with fsspec.open(_lowerCAmelCase			, '''rb'''			, compression='''infer'''					) as f:
																__SCREAMING_SNAKE_CASE:				Optional[Any]           =   f.read()
										with fsspec.open(_lowerCAmelCase			, '''rb'''			, compression='''infer'''					) as f:
																__SCREAMING_SNAKE_CASE:				Any           =   f.read()
										assert exported_content == original_content



