

import enum
import os
from hashlib import shaaaa
from typing import Optional

from .. import config
from .logging import get_logger


A_ :				Any						 =    get_logger(__name__)



class 					_a		(enum.Enum       ):


					'''simple docstring'''






					UpperCAmelCase__:     Dict													= '''all_checks'''
					UpperCAmelCase__:     Optional[int]													= '''basic_checks'''
					UpperCAmelCase__:     Optional[Any]													= '''no_checks'''



class 					_a		(__magic_name__       ):


					'''simple docstring'''









class 					_a		(__magic_name__       ):


					'''simple docstring'''









class 					_a		(__magic_name__       ):


					'''simple docstring'''









class 					_a		(__magic_name__       ):


					'''simple docstring'''









def 		UpperCamelCase      (lowercase_:				Optional[dict]					,   lowercase_:				dict					,   lowercase_:				List[Any]=None      )  ->   Dict:
	if expected_checksums is None:
		logger.info("""Unable to verify checksums."""      )
		return
	if len(set(lowercase_      ) - set(lowercase_      )      ) > 0:
		raise ExpectedMoreDownloadedFiles(str(set(lowercase_      ) - set(lowercase_      )      )      )
	if len(set(lowercase_      ) - set(lowercase_      )      ) > 0:
		raise UnexpectedDownloadedFile(str(set(lowercase_      ) - set(lowercase_      )      )      )
	A__							:					List[str]       =	[url for url in expected_checksums if expected_checksums[url] != recorded_checksums[url]]
	A__							:					Union[str, Any]       =	""" for """ + verification_name if verification_name is not None else """"""
	if len(lowercase_      ) > 0:
		raise NonMatchingChecksumError(
		    f"""Checksums didn't match{for_verification_name}:\n"""
		    f"""{bad_urls}\n"""
		    """Set `verification_mode='no_checks'` to skip checksums verification and ignore this error"""      )
	logger.info("""All the checksums matched successfully""" + for_verification_name      )



class 					_a		(__magic_name__       ):


					'''simple docstring'''









class 					_a		(__magic_name__       ):


					'''simple docstring'''









class 					_a		(__magic_name__       ):


					'''simple docstring'''









class 					_a		(__magic_name__       ):


					'''simple docstring'''









def 		UpperCamelCase      (lowercase_:				Optional[dict]					,   lowercase_:				dict      )  ->   int:
	if expected_splits is None:
		logger.info("""Unable to verify splits sizes."""      )
		return
	if len(set(lowercase_      ) - set(lowercase_      )      ) > 0:
		raise ExpectedMoreSplits(str(set(lowercase_      ) - set(lowercase_      )      )      )
	if len(set(lowercase_      ) - set(lowercase_      )      ) > 0:
		raise UnexpectedSplits(str(set(lowercase_      ) - set(lowercase_      )      )      )
	A__							:					int       =	[
	    {"""expected""": expected_splits[name], """recorded""": recorded_splits[name]}
	    for name in expected_splits
	    if expected_splits[name].num_examples != recorded_splits[name].num_examples
	]
	if len(lowercase_      ) > 0:
		raise NonMatchingSplitsSizesError(str(lowercase_      )      )
	logger.info("""All the splits matched successfully."""      )

def 		UpperCamelCase      (lowercase_:				str					,   lowercase_:				bool = True      )  ->   dict:
	if record_checksum:
		A__							:					List[Any]       =	shaaaa()
		with open(lowercase_					,   """rb"""      ) as f:
			for chunk in iter(lambda: f.read(1 << 20      )					,   b""""""      ):
				m.update(lowercase_      )
			A__							:					Any       =	m.hexdigest()
	else:
		A__							:					Union[str, Any]       =	None
	return {"num_bytes": os.path.getsize(lowercase_      ), "checksum": checksum}


def 		UpperCamelCase      (lowercase_:				Optional[Any]      )  ->   List[str]:
	if dataset_size and config.IN_MEMORY_MAX_SIZE:
		return dataset_size < config.IN_MEMORY_MAX_SIZE
	else:
		return False






import os
import tempfile
from functools import partial
from unittest import TestCase
from unittest.mock import patch

import datasets
import datasets.config

from .utils import require_beam



class 					_a		(datasets.BeamBasedBuilder       ):


					'''simple docstring'''






					def 						__A    (				self    ):
						return datasets.DatasetInfo(
						    features=datasets.Features({"""content""": datasets.Value("""string"""    )}    )					, supervised_keys=A__					, )






					def 						__A    (				self					, A__					, A__    ):
						return [datasets.SplitGenerator(name=datasets.Split.TRAIN					, gen_kwargs={"""examples""": get_test_dummy_examples()}    )]






					def 						__A    (				self					, A__					, A__    ):
						import apache_beam as beam

						return pipeline | "Load Examples" >> beam.Create(A__    )



class 					_a		(datasets.BeamBasedBuilder       ):


					'''simple docstring'''






					def 						__A    (				self    ):
						return datasets.DatasetInfo(
						    features=datasets.Features({"""a""": datasets.Sequence({"""b""": datasets.Value("""string"""    )}    )}    )					, supervised_keys=A__					, )






					def 						__A    (				self					, A__					, A__    ):
						return [
						    datasets.SplitGenerator(name=datasets.Split.TRAIN					, gen_kwargs={"""examples""": get_test_nested_examples()}    )
						]






					def 						__A    (				self					, A__					, A__    ):
						import apache_beam as beam

						return pipeline | "Load Examples" >> beam.Create(A__    )



def 		UpperCamelCase      ()  ->   Dict:
	return [(i, {"content": content}) for i, content in enumerate(["""foo""", """bar""", """foobar"""]      )]

def 		UpperCamelCase      ()  ->   Tuple:
	return [(i, {"a": {"b": [content]}}) for i, content in enumerate(["""foo""", """bar""", """foobar"""]      )]



class 					_a		(__magic_name__       ):


					'''simple docstring'''






					@require_beam
					def 						__A    (				self    ):
						A__							:					Dict       =	len(get_test_dummy_examples()    )
						with tempfile.TemporaryDirectory() as tmp_cache_dir:
							A__							:					int       =	DummyBeamDataset(cache_dir=A__					, beam_runner="""DirectRunner"""    )
							builder.download_and_prepare()
							self.assertTrue(
							    os.path.exists(
							        os.path.join(A__					, builder.name					, """default"""					, """0.0.0"""					, F"""{builder.name}-train.arrow"""    )    )    )
							self.assertDictEqual(builder.info.features					, datasets.Features({"""content""": datasets.Value("""string"""    )}    )    )
							A__							:					int       =	builder.as_dataset()
							self.assertEqual(dset["""train"""].num_rows					, A__    )
							self.assertEqual(dset["""train"""].info.splits["""train"""].num_examples					, A__    )
							self.assertDictEqual(dset["""train"""][0]					, get_test_dummy_examples()[0][1]    )
							self.assertDictEqual(
							    dset["""train"""][expected_num_examples - 1]					, get_test_dummy_examples()[expected_num_examples - 1][1]    )
							self.assertTrue(
							    os.path.exists(os.path.join(A__					, builder.name					, """default"""					, """0.0.0"""					, """dataset_info.json"""    )    )    )
							del dset






					@require_beam
					def 						__A    (				self    ):
						import apache_beam as beam

						A__							:					int       =	beam.io.parquetio.WriteToParquet

						A__							:					List[str]       =	len(get_test_dummy_examples()    )
						with tempfile.TemporaryDirectory() as tmp_cache_dir:
							A__							:					str       =	DummyBeamDataset(cache_dir=A__					, beam_runner="""DirectRunner"""    )
							with patch("""apache_beam.io.parquetio.WriteToParquet"""    ) as write_parquet_mock:
								A__							:					Optional[Any]       =	partial(A__					, num_shards=2    )
								builder.download_and_prepare()
							self.assertTrue(
							    os.path.exists(
							        os.path.join(
							            A__					, builder.name					, """default"""					, """0.0.0"""					, F"""{builder.name}-train-00000-of-00002.arrow"""    )    )    )
							self.assertTrue(
							    os.path.exists(
							        os.path.join(
							            A__					, builder.name					, """default"""					, """0.0.0"""					, F"""{builder.name}-train-00000-of-00002.arrow"""    )    )    )
							self.assertDictEqual(builder.info.features					, datasets.Features({"""content""": datasets.Value("""string"""    )}    )    )
							A__							:					Optional[int]       =	builder.as_dataset()
							self.assertEqual(dset["""train"""].num_rows					, A__    )
							self.assertEqual(dset["""train"""].info.splits["""train"""].num_examples					, A__    )
							# Order is not preserved when sharding, so we just check that all the elements are there
							self.assertListEqual(sorted(dset["""train"""]["""content"""]    )					, sorted(["""foo""", """bar""", """foobar"""]    )    )
							self.assertTrue(
							    os.path.exists(os.path.join(A__					, builder.name					, """default"""					, """0.0.0"""					, """dataset_info.json"""    )    )    )
							del dset






					@require_beam
					def 						__A    (				self    ):
						with tempfile.TemporaryDirectory() as tmp_cache_dir:
							A__							:					int       =	DummyBeamDataset(cache_dir=A__    )
							self.assertRaises(datasets.builder.MissingBeamOptions					, builder.download_and_prepare    )






					@require_beam
					def 						__A    (				self    ):
						A__							:					List[Any]       =	len(get_test_nested_examples()    )
						with tempfile.TemporaryDirectory() as tmp_cache_dir:
							A__							:					Optional[int]       =	NestedBeamDataset(cache_dir=A__					, beam_runner="""DirectRunner"""    )
							builder.download_and_prepare()
							self.assertTrue(
							    os.path.exists(
							        os.path.join(A__					, builder.name					, """default"""					, """0.0.0"""					, F"""{builder.name}-train.arrow"""    )    )    )
							self.assertDictEqual(
							    builder.info.features					, datasets.Features({"""a""": datasets.Sequence({"""b""": datasets.Value("""string"""    )}    )}    )    )
							A__							:					Optional[int]       =	builder.as_dataset()
							self.assertEqual(dset["""train"""].num_rows					, A__    )
							self.assertEqual(dset["""train"""].info.splits["""train"""].num_examples					, A__    )
							self.assertDictEqual(dset["""train"""][0]					, get_test_nested_examples()[0][1]    )
							self.assertDictEqual(
							    dset["""train"""][expected_num_examples - 1]					, get_test_nested_examples()[expected_num_examples - 1][1]    )
							self.assertTrue(
							    os.path.exists(os.path.join(A__					, builder.name					, """default"""					, """0.0.0"""					, """dataset_info.json"""    )    )    )
							del dset



