

'''simple docstring'''





import time
from dataclasses import dataclass
from multiprocessing import Pool
from unittest import TestCase
from unittest.mock import patch

import multiprocess
import numpy as np
import pytest

from datasets.utils.py_utils import (
    NestedDataStructure,
    asdict,
    iflatmap_unordered,
    map_nested,
    temp_seed,
    temporary_assignment,
    zip_dict,
)

from .utils import require_tf, require_torch







def   _UpperCamelCase	(		__UpperCamelCase	)       ->			Any:  # picklable for multiprocessing
					return x.sum()







def   _UpperCamelCase	(		__UpperCamelCase	)       ->			Union[str, Any]:  # picklable for multiprocessing
					return i + 1


@dataclass
class  UpperCAmelCase :

		'''simple docstring'''






		SCREAMING_SNAKE_CASE_           =					42
		SCREAMING_SNAKE_CASE_           =					42


class  UpperCAmelCase (							A__    ):

		'''simple docstring'''






		def 			UpperCamelCase(					self		)  ->     List[str]:






							'''simple docstring'''



							lowerCamelCase_  				=      {}
							lowerCamelCase_  				=      []
							lowerCamelCase_  				=      1
							lowerCamelCase_  				=      [1, 2]
							lowerCamelCase_  				=      {'a': 1, 'b': 2}
							lowerCamelCase_  				=      {'a': [1, 2], 'b': [3, 4]}
							lowerCamelCase_  				=      {'a': {'1': 1}, 'b': 2}
							lowerCamelCase_  				=      {'a': 1, 'b': 2, 'c': 3, 'd': 4}
							lowerCamelCase_  				=      {}
							lowerCamelCase_  				=      []
							lowerCamelCase_  				=      2
							lowerCamelCase_  				=      [2, 3]
							lowerCamelCase_  				=      {'a': 2, 'b': 3}
							lowerCamelCase_  				=      {'a': [2, 3], 'b': [4, 5]}
							lowerCamelCase_  				=      {'a': {'1': 2}, 'b': 3}
							lowerCamelCase_  				=      {'a': 2, 'b': 3, 'c': 4, 'd': 5}
							self.assertEqual(map_nested(_UpperCamelCase   ,			_UpperCamelCase		)   ,			_UpperCamelCase		)
							self.assertEqual(map_nested(_UpperCamelCase   ,			_UpperCamelCase		)   ,			_UpperCamelCase		)
							self.assertEqual(map_nested(_UpperCamelCase   ,			_UpperCamelCase		)   ,			_UpperCamelCase		)
							self.assertEqual(map_nested(_UpperCamelCase   ,			_UpperCamelCase		)   ,			_UpperCamelCase		)
							self.assertEqual(map_nested(_UpperCamelCase   ,			_UpperCamelCase		)   ,			_UpperCamelCase		)
							self.assertEqual(map_nested(_UpperCamelCase   ,			_UpperCamelCase		)   ,			_UpperCamelCase		)
							self.assertEqual(map_nested(_UpperCamelCase   ,			_UpperCamelCase		)   ,			_UpperCamelCase		)
							self.assertEqual(map_nested(_UpperCamelCase   ,			_UpperCamelCase		)   ,			_UpperCamelCase		)

							lowerCamelCase_  				=      2
							self.assertEqual(map_nested(_UpperCamelCase   ,			_UpperCamelCase   ,			num_proc=_UpperCamelCase		)   ,			_UpperCamelCase		)
							self.assertEqual(map_nested(_UpperCamelCase   ,			_UpperCamelCase   ,			num_proc=_UpperCamelCase		)   ,			_UpperCamelCase		)
							self.assertEqual(map_nested(_UpperCamelCase   ,			_UpperCamelCase   ,			num_proc=_UpperCamelCase		)   ,			_UpperCamelCase		)
							self.assertEqual(map_nested(_UpperCamelCase   ,			_UpperCamelCase   ,			num_proc=_UpperCamelCase		)   ,			_UpperCamelCase		)
							self.assertEqual(map_nested(_UpperCamelCase   ,			_UpperCamelCase   ,			num_proc=_UpperCamelCase		)   ,			_UpperCamelCase		)
							self.assertEqual(map_nested(_UpperCamelCase   ,			_UpperCamelCase   ,			num_proc=_UpperCamelCase		)   ,			_UpperCamelCase		)
							self.assertEqual(map_nested(_UpperCamelCase   ,			_UpperCamelCase   ,			num_proc=_UpperCamelCase		)   ,			_UpperCamelCase		)
							self.assertEqual(map_nested(_UpperCamelCase   ,			_UpperCamelCase   ,			num_proc=_UpperCamelCase		)   ,			_UpperCamelCase		)

							lowerCamelCase_  				=      {'a': np.eye(2		), 'b': np.zeros(3		), 'c': np.ones(2		)}
							lowerCamelCase_  				=      {'a': 2, 'b': 0, 'c': 2}
							lowerCamelCase_  				=      {
							    'a': np.eye(2		).astype(_UpperCamelCase		),
							    'b': np.zeros(3		).astype(_UpperCamelCase		),
							    'c': np.ones(2		).astype(_UpperCamelCase		),
							}
							self.assertEqual(map_nested(_UpperCamelCase   ,			_UpperCamelCase   ,			map_numpy=_UpperCamelCase		)   ,			_UpperCamelCase		)
							self.assertEqual(
							    {k: v.tolist() for k, v in map_nested(_UpperCamelCase   ,			_UpperCamelCase   ,			map_numpy=_UpperCamelCase		).items()}   ,			{k: v.tolist() for k, v in expected_map_nested_sna_int.items()}   ,			)
							self.assertEqual(map_nested(_UpperCamelCase   ,			_UpperCamelCase   ,			map_numpy=_UpperCamelCase   ,			num_proc=_UpperCamelCase		)   ,			_UpperCamelCase		)
							self.assertEqual(
							    {k: v.tolist() for k, v in map_nested(_UpperCamelCase   ,			_UpperCamelCase   ,			map_numpy=_UpperCamelCase   ,			num_proc=_UpperCamelCase		).items()}   ,			{k: v.tolist() for k, v in expected_map_nested_sna_int.items()}   ,			)
							with self.assertRaises(_UpperCamelCase		):  # can't pickle a local lambda
												map_nested(lambda SCREAMING_SNAKE_CASE_		: x + 1   ,			_UpperCamelCase   ,			num_proc=_UpperCamelCase		)
		def 			UpperCamelCase(					self		)  ->     Optional[Any]:






							'''simple docstring'''



							lowerCamelCase_  				=      {'a': 1, 'b': 2}
							lowerCamelCase_  				=      {'a': 3, 'b': 4}
							lowerCamelCase_  				=      {'a': 5, 'b': 6}
							lowerCamelCase_  				=      sorted([('a', (1, 3, 5)), ('b', (2, 4, 6))]		)
							self.assertEqual(sorted(zip_dict(_UpperCamelCase   ,			_UpperCamelCase   ,			_UpperCamelCase		)		)   ,			_UpperCamelCase		)






		def 			UpperCamelCase(					self		)  ->     str:






							'''simple docstring'''



							class  UpperCAmelCase :

									'''simple docstring'''






									SCREAMING_SNAKE_CASE_           =					"""bar"""

							lowerCamelCase_  				=      Foo()
							self.assertEqual(foo.my_attr   ,			'bar'		)
							with temporary_assignment(_UpperCamelCase   ,			'my_attr'   ,			'BAR'		):
												self.assertEqual(foo.my_attr   ,			'BAR'		)
							self.assertEqual(foo.my_attr   ,			'bar'		)







@pytest.mark.parametrize(
    'iterable_length, num_proc, expected_num_proc'    ,[
        (1, None, 1),
        (1, 1, 1),
        (2, None, 1),
        (2, 1, 1),
        (2, 2, 1),
        (2, 3, 1),
        (3, 2, 1),
        (16, 16, 16),
        (16, 17, 16),
        (17, 16, 16),
    ]    ,)
def   _UpperCamelCase	(		__UpperCamelCase    ,__UpperCamelCase    ,__UpperCamelCase	)       ->			Tuple:







					with patch('datasets.utils.py_utils._single_map_nested'	) as mock_single_map_nested, patch(
					    'datasets.parallel.parallel.Pool'	) as mock_multiprocessing_pool:
										lowerCamelCase_  				=      {f'''{i}''': i for i in range(SCREAMING_SNAKE_CASE_	)}
										lowerCamelCase_  				=      map_nested(lambda __UpperCamelCase	: x + 10    ,SCREAMING_SNAKE_CASE_    ,num_proc=SCREAMING_SNAKE_CASE_    ,parallel_min_length=16	)
										if expected_num_proc == 1:
															assert mock_single_map_nested.called
															assert not mock_multiprocessing_pool.called
										else:
															assert not mock_single_map_nested.called
															assert mock_multiprocessing_pool.called
															assert mock_multiprocessing_pool.call_args[0][0] == expected_num_proc


class  UpperCAmelCase (							A__    ):

		'''simple docstring'''






		@require_tf
		def 			UpperCamelCase(					self		)  ->     Any:






							'''simple docstring'''



							import tensorflow as tf
							from tensorflow.keras import layers

							lowerCamelCase_  				=      layers.Dense(2		)

							def gen_random_output():
												lowerCamelCase_  				=      tf.random.uniform((1, 3)		)
												return model(_UpperCamelCase		).numpy()

							with temp_seed(42   ,			set_tensorflow=_UpperCamelCase		):
												lowerCamelCase_  				=      gen_random_output()
							with temp_seed(42   ,			set_tensorflow=_UpperCamelCase		):
												lowerCamelCase_  				=      gen_random_output()
							lowerCamelCase_  				=      gen_random_output()

							np.testing.assert_equal(_UpperCamelCase   ,			_UpperCamelCase		)
							self.assertGreater(np.abs(outa - outa		).sum()   ,			0		)
		@require_torch
		def 			UpperCamelCase(					self		)  ->     Union[str, Any]:






							'''simple docstring'''



							import torch

							def gen_random_output():
												lowerCamelCase_  				=      torch.nn.Linear(3   ,			2		)
												lowerCamelCase_  				=      torch.rand(1   ,			3		)
												return model(_UpperCamelCase		).detach().numpy()

							with temp_seed(42   ,			set_pytorch=_UpperCamelCase		):
												lowerCamelCase_  				=      gen_random_output()
							with temp_seed(42   ,			set_pytorch=_UpperCamelCase		):
												lowerCamelCase_  				=      gen_random_output()
							lowerCamelCase_  				=      gen_random_output()

							np.testing.assert_equal(_UpperCamelCase   ,			_UpperCamelCase		)
							self.assertGreater(np.abs(outa - outa		).sum()   ,			0		)






		def 			UpperCamelCase(					self		)  ->     Dict:






							'''simple docstring'''



							def gen_random_output():
												return np.random.rand(1   ,			3		)

							with temp_seed(42		):
												lowerCamelCase_  				=      gen_random_output()
							with temp_seed(42		):
												lowerCamelCase_  				=      gen_random_output()
							lowerCamelCase_  				=      gen_random_output()

							np.testing.assert_equal(_UpperCamelCase   ,			_UpperCamelCase		)
							self.assertGreater(np.abs(outa - outa		).sum()   ,			0		)







@pytest.mark.parametrize('input_data'    ,[{}]	)
def   _UpperCamelCase	(		__UpperCamelCase	)       ->			Dict:
					lowerCamelCase_  				=      NestedDataStructure(SCREAMING_SNAKE_CASE_	).data
					assert output_data == input_data







@pytest.mark.parametrize(
    'data, expected_output'    ,[
        ({}, []),
        ([], []),
        ('foo', ['foo']),
        (['foo', 'bar'], ['foo', 'bar']),
        ([['foo', 'bar']], ['foo', 'bar']),
        ([[['foo'], ['bar']]], ['foo', 'bar']),
        ([[['foo'], 'bar']], ['foo', 'bar']),
        ({'a': 1, 'b': 2}, [1, 2]),
        ({'a': [1, 2], 'b': [3, 4]}, [1, 2, 3, 4]),
        ({'a': [[1, 2]], 'b': [[3, 4]]}, [1, 2, 3, 4]),
        ({'a': [[1, 2]], 'b': [3, 4]}, [1, 2, 3, 4]),
        ({'a': [[[1], [2]]], 'b': [[[3], [4]]]}, [1, 2, 3, 4]),
        ({'a': [[[1], [2]]], 'b': [[3, 4]]}, [1, 2, 3, 4]),
        ({'a': [[[1], [2]]], 'b': [3, 4]}, [1, 2, 3, 4]),
        ({'a': [[[1], [2]]], 'b': [3, [4]]}, [1, 2, 3, 4]),
        ({'a': {'1': 1}, 'b': 2}, [1, 2]),
        ({'a': {'1': [1]}, 'b': 2}, [1, 2]),
        ({'a': {'1': [1]}, 'b': [2]}, [1, 2]),
    ]    ,)
def   _UpperCamelCase	(		__UpperCamelCase    ,__UpperCamelCase	)       ->			Optional[Any]:
					lowerCamelCase_  				=      NestedDataStructure(SCREAMING_SNAKE_CASE_	).flatten()
					assert output == expected_output







def   _UpperCamelCase	(		)       ->			int:
					lowerCamelCase_  				=      A(x=1    ,y='foobar'	)
					lowerCamelCase_  				=      {'x': 1, 'y': 'foobar'}
					assert asdict(SCREAMING_SNAKE_CASE_	) == expected_output

					lowerCamelCase_  				=      {'a': {'b': A(x=10    ,y='foo'	)}, 'c': [A(x=20    ,y='bar'	)]}
					lowerCamelCase_  				=      {'a': {'b': {'x': 10, 'y': 'foo'}}, 'c': [{'x': 20, 'y': 'bar'}]}
					assert asdict(SCREAMING_SNAKE_CASE_	) == expected_output

					with pytest.raises(SCREAMING_SNAKE_CASE_	):
										asdict([1, A(x=10    ,y='foo'	)]	)







def   _UpperCamelCase	(		__UpperCamelCase	)       ->			int:
					return text.split()







def   _UpperCamelCase	(		__UpperCamelCase	)       ->			Any:
					yield (time.time(), content)
					time.sleep(2	)
					yield (time.time(), content)







def   _UpperCamelCase	(		)       ->			str:







					with Pool(2	) as pool:
										lowerCamelCase_  				=      list(iflatmap_unordered(SCREAMING_SNAKE_CASE_    ,_split_text    ,kwargs_iterable=[{'text': 'hello there'}] * 10	)	)
										assert out.count('hello'	) == 10
										assert out.count('there'	) == 10
										assert len(SCREAMING_SNAKE_CASE_	) == 20

					# check multiprocess from pathos (uses dill for pickling)
					with multiprocess.Pool(2	) as pool:
										lowerCamelCase_  				=      list(iflatmap_unordered(SCREAMING_SNAKE_CASE_    ,_split_text    ,kwargs_iterable=[{'text': 'hello there'}] * 10	)	)
										assert out.count('hello'	) == 10
										assert out.count('there'	) == 10
										assert len(SCREAMING_SNAKE_CASE_	) == 20

					# check that we get items as fast as possible
					with Pool(2	) as pool:
										lowerCamelCase_  				=      []
										for yield_time, content in iflatmap_unordered(
										    SCREAMING_SNAKE_CASE_    ,_aseconds_generator_of_aitems_with_timing    ,kwargs_iterable=[{'content': 'a'}, {'content': 'b'}]	):
															assert yield_time < time.time() + 0.1, "we should each item directly after it was yielded"
															out.append(SCREAMING_SNAKE_CASE_	)
										assert out.count('a'	) == 2
										assert out.count('b'	) == 2
										assert len(SCREAMING_SNAKE_CASE_	) == 4










from ..utils import DummyObject, requires_backends




class 			__UpperCamelCase    (							metaclass=A__    ):
     __A	:  str							  =   ["""torch""", """scipy"""]





     def __init__(     self				,							*_UpperCamelCase				,							**_UpperCamelCase					):
        requires_backends(self				,							['''torch''', '''scipy''']					)





     @classmethod
     def 		UpperCamelCase(     cls				,							*_UpperCamelCase				,							**_UpperCamelCase					):
        requires_backends(cls				,							['''torch''', '''scipy''']					)





     @classmethod
     def 		UpperCamelCase(     cls				,							*_UpperCamelCase				,							**_UpperCamelCase					):
        requires_backends(cls				,							['''torch''', '''scipy''']					)