
"""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 							a__ (				lowerCAmelCase	:       List[Any]		):  # picklable for multiprocessing






     '''simple docstring'''
     return x.sum()



def 							a__ (				lowerCAmelCase	:       Optional[int]		):  # picklable for multiprocessing






     '''simple docstring'''
     return i + 1







@dataclass
class   _lowercase						:


     '''simple docstring'''





     _A								=							42
     _A								=							42







class   _lowercase						(       lowerCAmelCase_				):


     '''simple docstring'''





     def  lowerCAmelCase__       (      self  )->					Dict:
          UpperCAmelCase__							: Union[str, Any] 							=						{}
          UpperCAmelCase__							: Dict 							=						[]
          UpperCAmelCase__							: Dict 							=						1
          UpperCAmelCase__							: List[Any] 							=						[1, 2]
          UpperCAmelCase__							: str 							=						{"a": 1, "b": 2}
          UpperCAmelCase__							: Any 							=						{"a": [1, 2], "b": [3, 4]}
          UpperCAmelCase__							: str 							=						{"a": {"1": 1}, "b": 2}
          UpperCAmelCase__							: Union[str, Any] 							=						{"a": 1, "b": 2, "c": 3, "d": 4}
          UpperCAmelCase__							: List[Any] 							=						{}
          UpperCAmelCase__							: Optional[Any] 							=						[]
          UpperCAmelCase__							: Optional[Any] 							=						2
          UpperCAmelCase__							: Optional[Any] 							=						[2, 3]
          UpperCAmelCase__							: Optional[Any] 							=						{"a": 2, "b": 3}
          UpperCAmelCase__							: Any 							=						{"a": [2, 3], "b": [4, 5]}
          UpperCAmelCase__							: int 							=						{"a": {"1": 2}, "b": 3}
          UpperCAmelCase__							: Optional[Any] 							=						{"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  )

          UpperCAmelCase__							: Optional[Any] 							=						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  )

          UpperCAmelCase__							: List[str] 							=						{"a": np.eye(2  ), "b": np.zeros(3  ), "c": np.ones(2  )}
          UpperCAmelCase__							: Union[str, Any] 							=						{"a": 2, "b": 0, "c": 2}
          UpperCAmelCase__							: Optional[Any] 							=						{
              "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 __UpperCamelCase  : x + 1 ,			__UpperCamelCase ,			num_proc=__UpperCamelCase  )



     def  lowerCAmelCase__       (      self  )->					List[Any]:
          UpperCAmelCase__							: int 							=						{"a": 1, "b": 2}
          UpperCAmelCase__							: int 							=						{"a": 3, "b": 4}
          UpperCAmelCase__							: List[str] 							=						{"a": 5, "b": 6}
          UpperCAmelCase__							: str 							=						sorted([("a", (1, 3, 5)), ("b", (2, 4, 6))]  )
          self.assertEqual(sorted(zip_dict(__UpperCamelCase ,			__UpperCamelCase ,			__UpperCamelCase  )  ) ,			__UpperCamelCase  )



     def  lowerCAmelCase__       (      self  )->					List[Any]:







          class   _lowercase						:


               '''simple docstring'''





               _A								=							'bar'





          UpperCAmelCase__							: List[Any] 							=						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 							a__ (				lowerCAmelCase	:       List[str]       ,    lowerCAmelCase	:       Union[str, Any]       ,    lowerCAmelCase	:       Dict		):






     '''simple docstring'''
     with patch("datasets.utils.py_utils._single_map_nested"		) as mock_single_map_nested, patch(
         "datasets.parallel.parallel.Pool"		) as mock_multiprocessing_pool:
          UpperCAmelCase__							: int 							=						{F"{i}": i for i in range(lowerCAmelCase		)}
          UpperCAmelCase__							: Dict 							=						map_nested(lambda lowerCAmelCase		: x + 10       ,    lowerCAmelCase       ,    num_proc=lowerCAmelCase       ,    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   _lowercase						(       lowerCAmelCase_				):


     '''simple docstring'''





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

          UpperCAmelCase__							: Dict 							=						layers.Dense(2  )

          def gen_random_output():
               UpperCAmelCase__							: str 							=						tf.random.uniform((1, 3)  )
               return model(__UpperCamelCase  ).numpy()

          with temp_seed(42 ,			set_tensorflow=__UpperCamelCase  ):
               UpperCAmelCase__							: List[Any] 							=						gen_random_output()
          with temp_seed(42 ,			set_tensorflow=__UpperCamelCase  ):
               UpperCAmelCase__							: Optional[int] 							=						gen_random_output()
          UpperCAmelCase__							: int 							=						gen_random_output()

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



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

          def gen_random_output():
               UpperCAmelCase__							: int 							=						torch.nn.Linear(3 ,			2  )
               UpperCAmelCase__							: Dict 							=						torch.rand(1 ,			3  )
               return model(__UpperCamelCase  ).detach().numpy()

          with temp_seed(42 ,			set_pytorch=__UpperCamelCase  ):
               UpperCAmelCase__							: List[str] 							=						gen_random_output()
          with temp_seed(42 ,			set_pytorch=__UpperCamelCase  ):
               UpperCAmelCase__							: Optional[Any] 							=						gen_random_output()
          UpperCAmelCase__							: List[str] 							=						gen_random_output()

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



     def  lowerCAmelCase__       (      self  )->					str:
          def gen_random_output():
               return np.random.rand(1 ,			3  )

          with temp_seed(42  ):
               UpperCAmelCase__							: Tuple 							=						gen_random_output()
          with temp_seed(42  ):
               UpperCAmelCase__							: Optional[int] 							=						gen_random_output()
          UpperCAmelCase__							: int 							=						gen_random_output()

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





@pytest.mark.parametrize("input_data"       ,    [{}]		)
def 							a__ (				lowerCAmelCase	:       List[Any]		):






     '''simple docstring'''
     UpperCAmelCase__							: Dict 							=						NestedDataStructure(lowerCAmelCase		).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 							a__ (				lowerCAmelCase	:       str       ,    lowerCAmelCase	:       Union[str, Any]		):






     '''simple docstring'''
     UpperCAmelCase__							: int 							=						NestedDataStructure(lowerCAmelCase		).flatten()
     assert output == expected_output



def 							a__ (				):






     '''simple docstring'''
     UpperCAmelCase__							: Tuple 							=						A(x=1       ,    y="foobar"		)
     UpperCAmelCase__							: Dict 							=						{"x": 1, "y": "foobar"}
     assert asdict(lowerCAmelCase		) == expected_output

     UpperCAmelCase__							: Any 							=						{"a": {"b": A(x=10       ,    y="foo"		)}, "c": [A(x=20       ,    y="bar"		)]}
     UpperCAmelCase__							: Optional[int] 							=						{"a": {"b": {"x": 10, "y": "foo"}}, "c": [{"x": 20, "y": "bar"}]}
     assert asdict(lowerCAmelCase		) == expected_output

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



def 							a__ (				lowerCAmelCase	:       str		):






     '''simple docstring'''
     return text.split()



def 							a__ (				lowerCAmelCase	:       Any		):






     '''simple docstring'''
     yield (time.time(), content)
     time.sleep(2		)
     yield (time.time(), content)



def 							a__ (				):






     '''simple docstring'''
     with Pool(2		) as pool:
          UpperCAmelCase__							: Dict 							=						list(iflatmap_unordered(lowerCAmelCase       ,    _split_text       ,    kwargs_iterable=[{"text": "hello there"}] * 10		)		)
          assert out.count("hello"		) == 10
          assert out.count("there"		) == 10
          assert len(lowerCAmelCase		) == 20

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

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





"""simple docstring"""




import timeit

import numpy as np

import datasets
from datasets.arrow_writer import ArrowWriter
from datasets.features.features import _ArrayXD



def 							a__ (				lowerCAmelCase	:       List[str]		):






     '''simple docstring'''
     def wrapper(*lowerCAmelCase	:       Any       ,    **lowerCAmelCase	:       Tuple		):
          UpperCAmelCase__							: Optional[int] 							=						timeit.default_timer()
          UpperCAmelCase__							: int 							=						func(*lowerCAmelCase       ,    **lowerCAmelCase		)
          UpperCAmelCase__							: List[Any] 							=						timeit.default_timer() - starttime
          return delta

     UpperCAmelCase__							: int 							=						func.__name__

     return wrapper



def 							a__ (				lowerCAmelCase	:       dict       ,    lowerCAmelCase	:       Optional[int]=100       ,    lowerCAmelCase	:       List[str]=None		):






     '''simple docstring'''
     UpperCAmelCase__							: str 							=						[]
     UpperCAmelCase__							: Optional[Any] 							=						seq_shapes or {}
     for i in range(lowerCAmelCase		):
          UpperCAmelCase__							: int 							=						{}
          for col_id, (k, v) in enumerate(features.items()		):
               if isinstance(lowerCAmelCase       ,    _ArrayXD		):
                    UpperCAmelCase__							: List[str] 							=						np.random.rand(*v.shape		).astype(v.dtype		)
               elif isinstance(lowerCAmelCase       ,    datasets.Value		):
                    if v.dtype == "string":
                         UpperCAmelCase__							: Dict 							=						"The small grey turtle was surprisingly fast when challenged."
                    else:
                         UpperCAmelCase__							: str 							=						np.random.randint(10       ,    size=1		).astype(v.dtype		).item()
               elif isinstance(lowerCAmelCase       ,    datasets.Sequence		):
                    while isinstance(lowerCAmelCase       ,    datasets.Sequence		):
                         UpperCAmelCase__							: List[str] 							=						v.feature
                    UpperCAmelCase__							: Optional[int] 							=						seq_shapes[k]
                    UpperCAmelCase__							: Optional[int] 							=						np.random.rand(*lowerCAmelCase		).astype(v.dtype		)
               UpperCAmelCase__							: Union[str, Any] 							=						data

          dummy_data.append((i, example)		)

     return dummy_data



def 							a__ (				lowerCAmelCase	:       List[str]       ,    lowerCAmelCase	:       Tuple       ,    lowerCAmelCase	:       List[str]=100       ,    lowerCAmelCase	:       Optional[int]=None		):






     '''simple docstring'''
     UpperCAmelCase__							: int 							=						generate_examples(lowerCAmelCase       ,    num_examples=lowerCAmelCase       ,    seq_shapes=lowerCAmelCase		)

     with ArrowWriter(features=lowerCAmelCase       ,    path=lowerCAmelCase		) as writer:
          for key, record in dummy_data:
               UpperCAmelCase__							: List[Any] 							=						features.encode_example(lowerCAmelCase		)
               writer.write(lowerCAmelCase		)

          UpperCAmelCase__							,	UpperCAmelCase__							: Optional[Any] 							=						writer.finalize()

     if not num_final_examples == num_examples:
          raise ValueError(
              F"Error writing the dataset, wrote {num_final_examples} examples but should have written {num_examples}."		)

     UpperCAmelCase__							: Optional[int] 							=						datasets.Dataset.from_file(filename=lowerCAmelCase       ,    info=datasets.DatasetInfo(features=lowerCAmelCase		)		)

     return dataset



