



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   ):



       def 		_lowerCamelCase      (		self							:					Optional[int]):


           '''simple docstring'''

           return datasets.DatasetInfo(
               features=datasets.Features({'''content''': datasets.Value('''string''')})					,		supervised_keys=__SCREAMING_SNAKE_CASE					,		)



       def 		_lowerCamelCase      (		self							:					Any					,		__SCREAMING_SNAKE_CASE							:					List[Any]					,		__SCREAMING_SNAKE_CASE							:					Any):


           '''simple docstring'''

           return [datasets.SplitGenerator(name=datasets.Split.TRAIN					,		gen_kwargs={'''examples''': get_test_dummy_examples()})]





       def 		_lowerCamelCase      (		self							:					List[str]					,		__SCREAMING_SNAKE_CASE							:					Any					,		__SCREAMING_SNAKE_CASE							:					Any):


           '''simple docstring'''

           import apache_beam as beam

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


class    _A	(     datasets.BeamBasedBuilder   ):



       def 		_lowerCamelCase      (		self							:					int):


           '''simple docstring'''

           return datasets.DatasetInfo(
               features=datasets.Features({'''a''': datasets.Sequence({'''b''': datasets.Value('''string''')})})					,		supervised_keys=__SCREAMING_SNAKE_CASE					,		)



       def 		_lowerCamelCase      (		self							:					Tuple					,		__SCREAMING_SNAKE_CASE							:					Any					,		__SCREAMING_SNAKE_CASE							:					Union[str, Any]):


           '''simple docstring'''

           return [
               datasets.SplitGenerator(name=datasets.Split.TRAIN					,		gen_kwargs={'''examples''': get_test_nested_examples()})
           ]





       def 		_lowerCamelCase      (		self							:					Optional[Any]					,		__SCREAMING_SNAKE_CASE							:					Dict					,		__SCREAMING_SNAKE_CASE							:					str):


           '''simple docstring'''

           import apache_beam as beam

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





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





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

class    _A	(     __UpperCAmelCase   ):



       @require_beam
       def 		_lowerCamelCase      (		self							:					str):


           '''simple docstring'''

           __a							=    len(get_test_dummy_examples())
           with tempfile.TemporaryDirectory() as tmp_cache_dir:
               __a							=    DummyBeamDataset(cache_dir=__SCREAMING_SNAKE_CASE					,		beam_runner='''DirectRunner''')
               builder.download_and_prepare()
               self.assertTrue(
                   os.path.exists(
                       os.path.join(__SCREAMING_SNAKE_CASE					,		builder.name					,		'''default'''					,		'''0.0.0'''					,		F'{builder.name}-train.arrow')))
               self.assertDictEqual(builder.info.features					,		datasets.Features({'''content''': datasets.Value('''string''')}))
               __a							=    builder.as_dataset()
               self.assertEqual(dset['''train'''].num_rows					,		__SCREAMING_SNAKE_CASE)
               self.assertEqual(dset['''train'''].info.splits['''train'''].num_examples					,		__SCREAMING_SNAKE_CASE)
               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(__SCREAMING_SNAKE_CASE					,		builder.name					,		'''default'''					,		'''0.0.0'''					,		'''dataset_info.json''')))
               del dset



       @require_beam
       def 		_lowerCamelCase      (		self							:					List[Any]):


           '''simple docstring'''

           import apache_beam as beam

           __a							=    beam.io.parquetio.WriteToParquet

           __a							=    len(get_test_dummy_examples())
           with tempfile.TemporaryDirectory() as tmp_cache_dir:
               __a							=    DummyBeamDataset(cache_dir=__SCREAMING_SNAKE_CASE					,		beam_runner='''DirectRunner''')
               with patch('''apache_beam.io.parquetio.WriteToParquet''') as write_parquet_mock:
                   __a							=    partial(__SCREAMING_SNAKE_CASE					,		num_shards=2)
                   builder.download_and_prepare()
               self.assertTrue(
                   os.path.exists(
                       os.path.join(
                           __SCREAMING_SNAKE_CASE					,		builder.name					,		'''default'''					,		'''0.0.0'''					,		F'{builder.name}-train-00000-of-00002.arrow')))
               self.assertTrue(
                   os.path.exists(
                       os.path.join(
                           __SCREAMING_SNAKE_CASE					,		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							=    builder.as_dataset()
               self.assertEqual(dset['''train'''].num_rows					,		__SCREAMING_SNAKE_CASE)
               self.assertEqual(dset['''train'''].info.splits['''train'''].num_examples					,		__SCREAMING_SNAKE_CASE)
               # 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(__SCREAMING_SNAKE_CASE					,		builder.name					,		'''default'''					,		'''0.0.0'''					,		'''dataset_info.json''')))
               del dset



       @require_beam
       def 		_lowerCamelCase      (		self							:					Union[str, Any]):


           '''simple docstring'''

           with tempfile.TemporaryDirectory() as tmp_cache_dir:
               __a							=    DummyBeamDataset(cache_dir=__SCREAMING_SNAKE_CASE)
               self.assertRaises(datasets.builder.MissingBeamOptions					,		builder.download_and_prepare)





       @require_beam
       def 		_lowerCamelCase      (		self							:					Tuple):


           '''simple docstring'''

           __a							=    len(get_test_nested_examples())
           with tempfile.TemporaryDirectory() as tmp_cache_dir:
               __a							=    NestedBeamDataset(cache_dir=__SCREAMING_SNAKE_CASE					,		beam_runner='''DirectRunner''')
               builder.download_and_prepare()
               self.assertTrue(
                   os.path.exists(
                       os.path.join(__SCREAMING_SNAKE_CASE					,		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							=    builder.as_dataset()
               self.assertEqual(dset['''train'''].num_rows					,		__SCREAMING_SNAKE_CASE)
               self.assertEqual(dset['''train'''].info.splits['''train'''].num_examples					,		__SCREAMING_SNAKE_CASE)
               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(__SCREAMING_SNAKE_CASE					,		builder.name					,		'''default'''					,		'''0.0.0'''					,		'''dataset_info.json''')))
               del dset










# Lint as: python3
import itertools
import os
import re


__snake_case						:int						=					re.compile(r'''([A-Z]+)([A-Z][a-z])''')
__snake_case						:List[Any]						=					re.compile(r'''([a-z\d])([A-Z])''')

__snake_case						:Dict						=					re.compile(r'''(?<!_)_(?!_)''')
__snake_case						:Optional[int]						=					re.compile(r'''(_{2,})''')

__snake_case						:Optional[Any]						=					r'''^\w+(\.\w+)*$'''

__snake_case						:Union[str, Any]						=					r'''<>:/\|?*'''





def 					__snake_case     (      _UpperCAmelCase    ):
    __a							=    _uppercase_uppercase_re.sub(R'''\1_\2'''       ,			_UpperCAmelCase    )
    __a							=    _lowercase_uppercase_re.sub(R'''\1_\2'''       ,			_UpperCAmelCase    )
    return name.lower()





def 					__snake_case     (      _UpperCAmelCase    ):
    __a							=    _single_underscore_re.split(_UpperCAmelCase    )
    __a							=    [_multiple_underscores_re.split(_UpperCAmelCase    ) for n in name]
    return "".join(n.capitalize() for n in itertools.chain.from_iterable(_UpperCAmelCase    ) if n != ''''''    )





def 					__snake_case     (      _UpperCAmelCase    ):
    if os.path.basename(_UpperCAmelCase    ) != name:
        raise ValueError(f'Should be a dataset name, not a path: {name}'    )
    return camelcase_to_snakecase(_UpperCAmelCase    )





def 					__snake_case     (      _UpperCAmelCase       ,			_UpperCAmelCase    ):
    if os.path.basename(_UpperCAmelCase    ) != name:
        raise ValueError(f'Should be a dataset name, not a path: {name}'    )
    if not re.match(_split_re       ,			_UpperCAmelCase    ):
        raise ValueError(f'Split name should match \'{_split_re}\'\' but got \'{split}\'.'    )
    return f'{filename_prefix_for_name(_UpperCAmelCase    )}-{split}'





def 					__snake_case     (      _UpperCAmelCase       ,			_UpperCAmelCase       ,			_UpperCAmelCase       ,			_UpperCAmelCase=None    ):
    __a							=    filename_prefix_for_split(_UpperCAmelCase       ,			_UpperCAmelCase    )
    if filetype_suffix:
        prefix += f'.{filetype_suffix}'
    __a							=    os.path.join(_UpperCAmelCase       ,			_UpperCAmelCase    )
    return f'{filepath}*'





def 					__snake_case     (      _UpperCAmelCase       ,			_UpperCAmelCase       ,			_UpperCAmelCase       ,			_UpperCAmelCase=None       ,			_UpperCAmelCase=None    ):
    __a							=    filename_prefix_for_split(_UpperCAmelCase       ,			_UpperCAmelCase    )
    __a							=    os.path.join(_UpperCAmelCase       ,			_UpperCAmelCase    )

    if shard_lengths:
        __a							=    len(_UpperCAmelCase    )
        __a							=    [f'{prefix}-{shard_id:05d}-of-{num_shards:05d}' for shard_id in range(_UpperCAmelCase    )]
        if filetype_suffix:
            __a							=    [filename + f'.{filetype_suffix}' for filename in filenames]
        return filenames
    else:
        __a							=    prefix
        if filetype_suffix:
            filename += f'.{filetype_suffix}'
        return [filename]





