



"""simple docstring"""


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  lowercase__						(		self     :    Any   )->List[str]:
							return datasets.DatasetInfo(
							    features=datasets.Features({'''content''': datasets.Value('''string'''   )}   )				,   supervised_keys=A__				,   )



			def  lowercase__						(		self     :    Optional[int]				,   __UpperCamelCase     :    Optional[int]				,   __UpperCamelCase     :    Optional[Any]   )->Any:
							return [datasets.SplitGenerator(name=datasets.Split.TRAIN				,   gen_kwargs={'''examples''': get_test_dummy_examples()}   )]



			def  lowercase__						(		self     :    Tuple				,   __UpperCamelCase     :    int				,   __UpperCamelCase     :    Dict   )->Tuple:
							import apache_beam as beam

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



class 			_a (					datasets.BeamBasedBuilder):
			"""simple docstring"""



			def  lowercase__						(		self     :    int   )->Union[str, Any]:
							return datasets.DatasetInfo(
							    features=datasets.Features({'''a''': datasets.Sequence({'''b''': datasets.Value('''string'''   )}   )}   )				,   supervised_keys=A__				,   )



			def  lowercase__						(		self     :    Optional[Any]				,   __UpperCamelCase     :    str				,   __UpperCamelCase     :    str   )->List[str]:
							return [
							    datasets.SplitGenerator(name=datasets.Split.TRAIN				,   gen_kwargs={'''examples''': get_test_nested_examples()}   )
							]



			def  lowercase__						(		self     :    List[str]				,   __UpperCamelCase     :    List[Any]				,   __UpperCamelCase     :    List[Any]   )->Any:
							import apache_beam as beam

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







def  lowercase	(      ):






				'''simple docstring'''



				return [(i, {"content": content}) for i, content in enumerate(['''foo''', '''bar''', '''foobar''']				)]







def  lowercase	(      ):






				'''simple docstring'''



				return [(i, {"a": {"b": [content]}}) for i, content in enumerate(['''foo''', '''bar''', '''foobar''']				)]
class 			_a (					lowerCAmelCase):
			"""simple docstring"""



			@require_beam
			def  lowercase__						(		self     :    Optional[int]   )->Dict:
							_UpperCAmelCase					=    len(get_test_dummy_examples()   )
							with tempfile.TemporaryDirectory() as tmp_cache_dir:
											_UpperCAmelCase					=    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'''   )}   )   )
											_UpperCAmelCase					=    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  lowercase__						(		self     :    List[str]   )->Dict:
							import apache_beam as beam

							_UpperCAmelCase					=    beam.io.parquetio.WriteToParquet

							_UpperCAmelCase					=    len(get_test_dummy_examples()   )
							with tempfile.TemporaryDirectory() as tmp_cache_dir:
											_UpperCAmelCase					=    DummyBeamDataset(cache_dir=A__				,   beam_runner='''DirectRunner'''   )
											with patch('''apache_beam.io.parquetio.WriteToParquet'''   ) as write_parquet_mock:
															_UpperCAmelCase					=    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'''   )}   )   )
											_UpperCAmelCase					=    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  lowercase__						(		self     :    str   )->Any:
							with tempfile.TemporaryDirectory() as tmp_cache_dir:
											_UpperCAmelCase					=    DummyBeamDataset(cache_dir=A__   )
											self.assertRaises(datasets.builder.MissingBeamOptions				,   builder.download_and_prepare   )



			@require_beam
			def  lowercase__						(		self     :    Optional[int]   )->Union[str, Any]:
							_UpperCAmelCase					=    len(get_test_nested_examples()   )
							with tempfile.TemporaryDirectory() as tmp_cache_dir:
											_UpperCAmelCase					=    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'''   )}   )}   )   )
											_UpperCAmelCase					=    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








"""simple docstring"""


from typing import Optional

from .. import Features, NamedSplit
from ..packaged_modules.text.text import Text
from ..utils.typing import NestedDataStructureLike, PathLike
from .abc import AbstractDatasetReader
class 			_a (					lowerCAmelCase):
			"""simple docstring"""



			def __init__(		self     :    List[str]				,   __UpperCamelCase     :    NestedDataStructureLike[PathLike]				,   __UpperCamelCase     :    Optional[NamedSplit] = None				,   __UpperCamelCase     :    Optional[Features] = None				,   __UpperCamelCase     :    str = None				,   __UpperCamelCase     :    bool = False				,   __UpperCamelCase     :    bool = False				,   __UpperCamelCase     :    Optional[int] = None				,   **__UpperCamelCase     :    Optional[Any]				,   )->Any:
							super().__init__(
							    __UpperCamelCase				,   split=__UpperCamelCase				,   features=__UpperCamelCase				,   cache_dir=__UpperCamelCase				,   keep_in_memory=__UpperCamelCase				,   streaming=__UpperCamelCase				,   num_proc=__UpperCamelCase				,   **__UpperCamelCase				,   )
							_UpperCAmelCase					=    path_or_paths if isinstance(__UpperCamelCase				,   __UpperCamelCase   ) else {self.split: path_or_paths}
							_UpperCAmelCase					=    Text(
							    cache_dir=__UpperCamelCase				,   data_files=__UpperCamelCase				,   features=__UpperCamelCase				,   **__UpperCamelCase				,   )



			def  lowercase__						(		self     :    Optional[Any]   )->str:
							# Build iterable dataset
							if self.streaming:
											_UpperCAmelCase					=    self.builder.as_streaming_dataset(split=self.split   )
							# Build regular (map-style) dataset
							else:
											_UpperCAmelCase					=    None
											_UpperCAmelCase					=    None
											_UpperCAmelCase					=    None
											_UpperCAmelCase					=    None

											self.builder.download_and_prepare(
											    download_config=__UpperCamelCase				,   download_mode=__UpperCamelCase				,   verification_mode=__UpperCamelCase				,   base_path=__UpperCamelCase				,   num_proc=self.num_proc				,   )
											_UpperCAmelCase					=    self.builder.as_dataset(
											    split=self.split				,   verification_mode=__UpperCamelCase				,   in_memory=self.keep_in_memory   )
							return dataset



