




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







							def      __lowerCAmelCase    (						self :						Optional[Any]      )    ->							Dict:


														"""simple docstring"""


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







							def      __lowerCAmelCase    (						self :						Any							,    lowerCAmelCase_ :						Tuple							,    lowerCAmelCase_ :						Optional[Any]      )    ->							Dict:


														"""simple docstring"""


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







							def      __lowerCAmelCase    (						self :						Tuple							,    lowerCAmelCase_ :						Dict							,    lowerCAmelCase_ :						Optional[Any]      )    ->							Optional[Any]:


														"""simple docstring"""


														import apache_beam as beam

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



class     A					(	datasets.BeamBasedBuilder					):







							def      __lowerCAmelCase    (						self :						int      )    ->							str:


														"""simple docstring"""


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







							def      __lowerCAmelCase    (						self :						Union[str, Any]							,    lowerCAmelCase_ :						Union[str, Any]							,    lowerCAmelCase_ :						List[Any]      )    ->							Any:


														"""simple docstring"""


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







							def      __lowerCAmelCase    (						self :						Optional[int]							,    lowerCAmelCase_ :						Optional[Any]							,    lowerCAmelCase_ :						Optional[int]      )    ->							Dict:


														"""simple docstring"""


														import apache_beam as beam

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



def 	snake_case_					():





							'''simple docstring'''


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





							'''simple docstring'''


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

class     A					(	_a					):







							@require_beam
							def      __lowerCAmelCase    (						self :						int      )    ->							Dict:


														"""simple docstring"""


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







							@require_beam
							def      __lowerCAmelCase    (						self :						int      )    ->							Dict:


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







							@require_beam
							def      __lowerCAmelCase    (						self :						Dict      )    ->							str:


														"""simple docstring"""


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







							@require_beam
							def      __lowerCAmelCase    (						self :						Optional[int]      )    ->							Any:


														"""simple docstring"""


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







'''simple docstring'''


import os

import pytest

from transformers.dynamic_module_utils import get_imports


_snake_case   :   List[str]     							=	'\nimport os\n'

_snake_case   :   Dict     							=	'\ndef foo():\n    import os\n    return False\n'

_snake_case   :   List[Any]     							=	'\ndef foo():\n    def bar():\n        if True:\n            import os\n        return False\n    return bar()\n'

_snake_case   :   Dict     							=	'\nimport os\n\ntry:\n    import bar\nexcept ImportError:\n    raise ValueError()\n'

_snake_case   :   Optional[Any]     							=	'\nimport os\n\ndef foo():\n    try:\n        import bar\n    except ImportError:\n        raise ValueError()\n'

_snake_case   :   Optional[Any]     							=	'\nimport os\n\ntry:\n    import bar\nexcept (ImportError, AttributeError):\n    raise ValueError()\n'

_snake_case   :   Any     							=	'\nimport os\n\ntry:\n    import bar\nexcept ImportError as e:\n    raise ValueError()\n'

_snake_case   :   List[str]     							=	'\nimport os\n\ntry:\n    import bar\nexcept:\n    raise ValueError()\n'

_snake_case   :   Dict     							=	'\nimport os\n\ntry:\n    import bar\n    import baz\nexcept ImportError:\n    raise ValueError()\n'

_snake_case   :   Any     							=	'\nimport os\n\ntry:\n    import bar\n    import baz\nexcept ImportError:\n    x = 1\n    raise ValueError()\n'

_snake_case   :   Union[str, Any]     							=	[
    TOP_LEVEL_IMPORT,
    IMPORT_IN_FUNCTION,
    DEEPLY_NESTED_IMPORT,
    TOP_LEVEL_TRY_IMPORT,
    GENERIC_EXCEPT_IMPORT,
    MULTILINE_TRY_IMPORT,
    MULTILINE_BOTH_IMPORT,
    MULTIPLE_EXCEPTS_IMPORT,
    EXCEPT_AS_IMPORT,
    TRY_IMPORT_IN_FUNCTION,
]
@pytest.mark.parametrize('''case'''				,		UpperCamelCase   )
def 	snake_case_					(UpperCamelCase :						str				,		UpperCamelCase :						List[Any]   ):





							'''simple docstring'''


							_a			  =       os.path.join(UpperCamelCase				,		'''test_file.py'''   )
							with open(UpperCamelCase				,		'''w'''   ) as _tmp_file:
														_tmp_file.write(UpperCamelCase   )

							_a			  =       get_imports(UpperCamelCase   )
							assert parsed_imports == ["os"]

