






'''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 (				__A				: Optional[Any]   ):  # picklable for multiprocessing
				return x.sum()





def 		_UpperCAmelCase (				__A				: Any   ):  # picklable for multiprocessing
				return i + 1





@dataclass
class 						SCREAMING_SNAKE_CASE       :
			snake_case__	=					42
			snake_case__	=					42





class 						SCREAMING_SNAKE_CASE       (       SCREAMING_SNAKE_CASE_							):






			def       SCREAMING_SNAKE_CASE   (			self   :      List[str]       )      ->      List[Any]:
							a_		:	Union[str, Any]									=   {}
							a_		:	int									=   []
							a_		:	List[str]									=   1
							a_		:	int									=   [1, 2]
							a_		:	int									=   {'''a''': 1, '''b''': 2}
							a_		:	List[str]									=   {'''a''': [1, 2], '''b''': [3, 4]}
							a_		:	Optional[Any]									=   {'''a''': {'''1''': 1}, '''b''': 2}
							a_		:	Dict									=   {'''a''': 1, '''b''': 2, '''c''': 3, '''d''': 4}
							a_		:	str									=   {}
							a_		:	Optional[int]									=   []
							a_		:	Union[str, Any]									=   2
							a_		:	Optional[int]									=   [2, 3]
							a_		:	Dict									=   {'''a''': 2, '''b''': 3}
							a_		:	Tuple									=   {'''a''': [2, 3], '''b''': [4, 5]}
							a_		:	List[Any]									=   {'''a''': {'''1''': 2}, '''b''': 3}
							a_		:	str									=   {'''a''': 2, '''b''': 3, '''c''': 4, '''d''': 5}
							self.assertEqual(map_nested(__SCREAMING_SNAKE_CASE	,    __SCREAMING_SNAKE_CASE       )	,    __SCREAMING_SNAKE_CASE       )
							self.assertEqual(map_nested(__SCREAMING_SNAKE_CASE	,    __SCREAMING_SNAKE_CASE       )	,    __SCREAMING_SNAKE_CASE       )
							self.assertEqual(map_nested(__SCREAMING_SNAKE_CASE	,    __SCREAMING_SNAKE_CASE       )	,    __SCREAMING_SNAKE_CASE       )
							self.assertEqual(map_nested(__SCREAMING_SNAKE_CASE	,    __SCREAMING_SNAKE_CASE       )	,    __SCREAMING_SNAKE_CASE       )
							self.assertEqual(map_nested(__SCREAMING_SNAKE_CASE	,    __SCREAMING_SNAKE_CASE       )	,    __SCREAMING_SNAKE_CASE       )
							self.assertEqual(map_nested(__SCREAMING_SNAKE_CASE	,    __SCREAMING_SNAKE_CASE       )	,    __SCREAMING_SNAKE_CASE       )
							self.assertEqual(map_nested(__SCREAMING_SNAKE_CASE	,    __SCREAMING_SNAKE_CASE       )	,    __SCREAMING_SNAKE_CASE       )
							self.assertEqual(map_nested(__SCREAMING_SNAKE_CASE	,    __SCREAMING_SNAKE_CASE       )	,    __SCREAMING_SNAKE_CASE       )

							a_		:	Tuple									=   2
							self.assertEqual(map_nested(__SCREAMING_SNAKE_CASE	,    __SCREAMING_SNAKE_CASE	,    num_proc=__SCREAMING_SNAKE_CASE       )	,    __SCREAMING_SNAKE_CASE       )
							self.assertEqual(map_nested(__SCREAMING_SNAKE_CASE	,    __SCREAMING_SNAKE_CASE	,    num_proc=__SCREAMING_SNAKE_CASE       )	,    __SCREAMING_SNAKE_CASE       )
							self.assertEqual(map_nested(__SCREAMING_SNAKE_CASE	,    __SCREAMING_SNAKE_CASE	,    num_proc=__SCREAMING_SNAKE_CASE       )	,    __SCREAMING_SNAKE_CASE       )
							self.assertEqual(map_nested(__SCREAMING_SNAKE_CASE	,    __SCREAMING_SNAKE_CASE	,    num_proc=__SCREAMING_SNAKE_CASE       )	,    __SCREAMING_SNAKE_CASE       )
							self.assertEqual(map_nested(__SCREAMING_SNAKE_CASE	,    __SCREAMING_SNAKE_CASE	,    num_proc=__SCREAMING_SNAKE_CASE       )	,    __SCREAMING_SNAKE_CASE       )
							self.assertEqual(map_nested(__SCREAMING_SNAKE_CASE	,    __SCREAMING_SNAKE_CASE	,    num_proc=__SCREAMING_SNAKE_CASE       )	,    __SCREAMING_SNAKE_CASE       )
							self.assertEqual(map_nested(__SCREAMING_SNAKE_CASE	,    __SCREAMING_SNAKE_CASE	,    num_proc=__SCREAMING_SNAKE_CASE       )	,    __SCREAMING_SNAKE_CASE       )
							self.assertEqual(map_nested(__SCREAMING_SNAKE_CASE	,    __SCREAMING_SNAKE_CASE	,    num_proc=__SCREAMING_SNAKE_CASE       )	,    __SCREAMING_SNAKE_CASE       )

							a_		:	Any									=   {'''a''': np.eye(2       ), '''b''': np.zeros(3       ), '''c''': np.ones(2       )}
							a_		:	Dict									=   {'''a''': 2, '''b''': 0, '''c''': 2}
							a_		:	List[Any]									=   {
							    '''a''': np.eye(2       ).astype(__SCREAMING_SNAKE_CASE       ),
							    '''b''': np.zeros(3       ).astype(__SCREAMING_SNAKE_CASE       ),
							    '''c''': np.ones(2       ).astype(__SCREAMING_SNAKE_CASE       ),
							}
							self.assertEqual(map_nested(__SCREAMING_SNAKE_CASE	,    __SCREAMING_SNAKE_CASE	,    map_numpy=__SCREAMING_SNAKE_CASE       )	,    __SCREAMING_SNAKE_CASE       )
							self.assertEqual(
							    {k: v.tolist() for k, v in map_nested(__SCREAMING_SNAKE_CASE	,    __SCREAMING_SNAKE_CASE	,    map_numpy=__SCREAMING_SNAKE_CASE       ).items()}	,    {k: v.tolist() for k, v in expected_map_nested_sna_int.items()}	,    )
							self.assertEqual(map_nested(__SCREAMING_SNAKE_CASE	,    __SCREAMING_SNAKE_CASE	,    map_numpy=__SCREAMING_SNAKE_CASE	,    num_proc=__SCREAMING_SNAKE_CASE       )	,    __SCREAMING_SNAKE_CASE       )
							self.assertEqual(
							    {k: v.tolist() for k, v in map_nested(__SCREAMING_SNAKE_CASE	,    __SCREAMING_SNAKE_CASE	,    map_numpy=__SCREAMING_SNAKE_CASE	,    num_proc=__SCREAMING_SNAKE_CASE       ).items()}	,    {k: v.tolist() for k, v in expected_map_nested_sna_int.items()}	,    )
							with self.assertRaises(__SCREAMING_SNAKE_CASE       ):  # can't pickle a local lambda
											map_nested(lambda __SCREAMING_SNAKE_CASE       : x + 1	,    __SCREAMING_SNAKE_CASE	,    num_proc=__SCREAMING_SNAKE_CASE       )






			def       SCREAMING_SNAKE_CASE   (			self   :      Optional[Any]       )      ->      Tuple:
							a_		:	List[str]									=   {'''a''': 1, '''b''': 2}
							a_		:	str									=   {'''a''': 3, '''b''': 4}
							a_		:	List[str]									=   {'''a''': 5, '''b''': 6}
							a_		:	List[Any]									=   sorted([('''a''', (1, 3, 5)), ('''b''', (2, 4, 6))]       )
							self.assertEqual(sorted(zip_dict(__SCREAMING_SNAKE_CASE	,    __SCREAMING_SNAKE_CASE	,    __SCREAMING_SNAKE_CASE       )       )	,    __SCREAMING_SNAKE_CASE       )







			def       SCREAMING_SNAKE_CASE   (			self   :      Optional[int]       )      ->      List[str]:





							class 						SCREAMING_SNAKE_CASE       :
										snake_case__	=					"bar"





							a_		:	Union[str, Any]									=   Foo()
							self.assertEqual(foo.my_attr	,    '''bar'''       )
							with temporary_assignment(__SCREAMING_SNAKE_CASE	,    '''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 (				__A				: Optional[int]    ,		__A				: Optional[Any]    ,		__A				: int   ):
				with patch('''datasets.utils.py_utils._single_map_nested'''   ) as mock_single_map_nested, patch(
				    '''datasets.parallel.parallel.Pool'''   ) as mock_multiprocessing_pool:
								a_		:	str									=   {f'{i}': i for i in range(__A   )}
								a_		:	Optional[int]									=   map_nested(lambda __A   : x + 10    ,		__A    ,		num_proc=__A    ,		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 						SCREAMING_SNAKE_CASE       (       SCREAMING_SNAKE_CASE_							):






			@require_tf
			def       SCREAMING_SNAKE_CASE   (			self   :      Any       )      ->      Any:
							import tensorflow as tf
							from tensorflow.keras import layers

							a_		:	int									=   layers.Dense(2       )

							def gen_random_output():
											a_		:	List[Any]									=   tf.random.uniform((1, 3)       )
											return model(__SCREAMING_SNAKE_CASE       ).numpy()

							with temp_seed(42	,    set_tensorflow=__SCREAMING_SNAKE_CASE       ):
											a_		:	Optional[Any]									=   gen_random_output()
							with temp_seed(42	,    set_tensorflow=__SCREAMING_SNAKE_CASE       ):
											a_		:	Optional[Any]									=   gen_random_output()
							a_		:	Any									=   gen_random_output()

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






			@require_torch
			def       SCREAMING_SNAKE_CASE   (			self   :      List[Any]       )      ->      Optional[Any]:
							import torch

							def gen_random_output():
											a_		:	List[str]									=   torch.nn.Linear(3	,    2       )
											a_		:	List[Any]									=   torch.rand(1	,    3       )
											return model(__SCREAMING_SNAKE_CASE       ).detach().numpy()

							with temp_seed(42	,    set_pytorch=__SCREAMING_SNAKE_CASE       ):
											a_		:	List[Any]									=   gen_random_output()
							with temp_seed(42	,    set_pytorch=__SCREAMING_SNAKE_CASE       ):
											a_		:	str									=   gen_random_output()
							a_		:	Any									=   gen_random_output()

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







			def       SCREAMING_SNAKE_CASE   (			self   :      int       )      ->      Dict:
							def gen_random_output():
											return np.random.rand(1	,    3       )

							with temp_seed(42       ):
											a_		:	Tuple									=   gen_random_output()
							with temp_seed(42       ):
											a_		:	List[Any]									=   gen_random_output()
							a_		:	str									=   gen_random_output()

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





@pytest.mark.parametrize('''input_data'''    ,		[{}]   )
def 		_UpperCAmelCase (				__A				: Any   ):
				a_		:	List[Any]									=   NestedDataStructure(__A   ).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 (				__A				: Optional[int]    ,		__A				: Optional[int]   ):
				a_		:	str									=   NestedDataStructure(__A   ).flatten()
				assert output == expected_output





def 		_UpperCAmelCase (				):
				a_		:	Tuple									=   A(x=1    ,		y='''foobar'''   )
				a_		:	Union[str, Any]									=   {'''x''': 1, '''y''': '''foobar'''}
				assert asdict(__A   ) == expected_output

				a_		:	Any									=   {'''a''': {'''b''': A(x=10    ,		y='''foo'''   )}, '''c''': [A(x=20    ,		y='''bar'''   )]}
				a_		:	List[Any]									=   {'''a''': {'''b''': {'''x''': 10, '''y''': '''foo'''}}, '''c''': [{'''x''': 20, '''y''': '''bar'''}]}
				assert asdict(__A   ) == expected_output

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





def 		_UpperCAmelCase (				__A				: str   ):
				return text.split()





def 		_UpperCAmelCase (				__A				: Tuple   ):
				yield (time.time(), content)
				time.sleep(2   )
				yield (time.time(), content)





def 		_UpperCAmelCase (				):
				with Pool(2   ) as pool:
								a_		:	Optional[int]									=   list(iflatmap_unordered(__A    ,		_split_text    ,		kwargs_iterable=[{'''text''': '''hello there'''}] * 10   )   )
								assert out.count('''hello'''   ) == 10
								assert out.count('''there'''   ) == 10
								assert len(__A   ) == 20

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

				# check that we get items as fast as possible
				with Pool(2   ) as pool:
								a_		:	str									=   []
								for yield_time, content in iflatmap_unordered(
								    __A    ,		_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(__A   )
								assert out.count('''a'''   ) == 2
								assert out.count('''b'''   ) == 2
								assert len(__A   ) == 4












'''simple docstring'''







from __future__ import annotations





def 		_UpperCAmelCase (				__A				: str    ,		__A				: str   ):
				a_		:	int									=   get_failure_array(__A   )

				# 2) Step through text searching for pattern
				a_					,      a_		:	Any									=   0, 0  # index into text, pattern
				while i < len(__A   ):
								if pattern[j] == text[i]:
												if j == (len(__A   ) - 1):
																return True
												j += 1

								# if this is a prefix in our pattern
								# just go back far enough to continue
								elif j > 0:
												a_		:	Any									=   failure[j - 1]
												continue
								i += 1
				return False





def 		_UpperCAmelCase (				__A				: str   ):
				a_		:	Optional[Any]									=   [0]
				a_		:	Any									=   0
				a_		:	int									=   1
				while j < len(__A   ):
								if pattern[i] == pattern[j]:
												i += 1
								elif i > 0:
												a_		:	List[Any]									=   failure[i - 1]
												continue
								j += 1
								failure.append(__A   )
				return failure


if __name__ == "__main__":
		# Test 1)
		__lowerCAmelCase									=   'abc1abc12'
		__lowerCAmelCase									=   'alskfjaldsabc1abc1abc12k23adsfabcabc'
		__lowerCAmelCase									=   'alskfjaldsk23adsfabcabc'
		assert kmp(pattern, texta) and not kmp(pattern, texta)

		# Test 2)
		__lowerCAmelCase									=   'ABABX'
		__lowerCAmelCase									=   'ABABZABABYABABX'
		assert kmp(pattern, text)

		# Test 3)
		__lowerCAmelCase									=   'AAAB'
		__lowerCAmelCase									=   'ABAAAAAB'
		assert kmp(pattern, text)

		# Test 4)
		__lowerCAmelCase									=   'abcdabcy'
		__lowerCAmelCase									=   'abcxabcdabxabcdabcdabcy'
		assert kmp(pattern, text)

		# Test 5)
		__lowerCAmelCase									=   'aabaabaaa'
		assert get_failure_array(pattern) == [0, 1, 0, 1, 2, 3, 4, 5, 2]




