
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	(      lowercase_   )		->	Tuple:  # picklable for multiprocessing


  '''simple docstring'''






  return x.sum()

def 		UpperCamelCase	(      lowercase_   )		->	str:  # picklable for multiprocessing


  '''simple docstring'''






  return i + 1



@dataclass
class       _snake_case   :
 __lowerCAmelCase  :     int
 __lowerCAmelCase  :     str



class       _snake_case   (	UpperCAmelCase_							):


 def      lowercase__     (					self):


   '''simple docstring'''




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

   lowercase__  :   Dict    		=      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_)

   lowercase__  :   Optional[int]    		=      {"""a""": np.eye(2), """b""": np.zeros(3), """c""": np.ones(2)}
   lowercase__  :   Union[str, Any]    		=      {"""a""": 2, """b""": 0, """c""": 2}
   lowercase__  :   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      lowercase__     (					self):


   '''simple docstring'''




   lowercase__  :   Union[str, Any]    		=      {"""a""": 1, """b""": 2}
   lowercase__  :   str    		=      {"""a""": 3, """b""": 4}
   lowercase__  :   Union[str, Any]    		=      {"""a""": 5, """b""": 6}
   lowercase__  :   Dict    		=      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      lowercase__     (					self):


   '''simple docstring'''




   class       _snake_case   :
    __lowerCAmelCase  :     Any        =			'bar'

   lowercase__  :   int    		=      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	(      lowercase_ , lowercase_ , lowercase_   )		->	int:


  '''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:
    lowercase__  :   Optional[int]    		=      {F'{i}': i for i in range(lowercase_   )}
    lowercase__  :   int    		=      map_nested(lambda lowercase_   : x + 10 , lowercase_ , num_proc=lowercase_ , 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       _snake_case   (	UpperCAmelCase_							):


 @require_tf
 def      lowercase__     (					self):


   '''simple docstring'''




   import tensorflow as tf
   from tensorflow.keras import layers

   lowercase__  :   int    		=      layers.Dense(2)

   def gen_random_output():
     lowercase__  :   Optional[int]    		=      tf.random.uniform((1, 3))
     return model(SCREAMING_SNAKE_CASE_).numpy()

   with temp_seed(42					,  set_tensorflow=SCREAMING_SNAKE_CASE_):
     lowercase__  :   Union[str, Any]    		=      gen_random_output()
   with temp_seed(42					,  set_tensorflow=SCREAMING_SNAKE_CASE_):
     lowercase__  :   Tuple    		=      gen_random_output()
   lowercase__  :   Tuple    		=      gen_random_output()

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


 @require_torch
 def      lowercase__     (					self):


   '''simple docstring'''




   import torch

   def gen_random_output():
     lowercase__  :   Union[str, Any]    		=      torch.nn.Linear(3					,  2)
     lowercase__  :   int    		=      torch.rand(1					,  3)
     return model(SCREAMING_SNAKE_CASE_).detach().numpy()

   with temp_seed(42					,  set_pytorch=SCREAMING_SNAKE_CASE_):
     lowercase__  :   str    		=      gen_random_output()
   with temp_seed(42					,  set_pytorch=SCREAMING_SNAKE_CASE_):
     lowercase__  :   List[Any]    		=      gen_random_output()
   lowercase__  :   List[Any]    		=      gen_random_output()

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



 def      lowercase__     (					self):


   '''simple docstring'''




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

   with temp_seed(42):
     lowercase__  :   Tuple    		=      gen_random_output()
   with temp_seed(42):
     lowercase__  :   Dict    		=      gen_random_output()
   lowercase__  :   List[Any]    		=      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	(      lowercase_   )		->	Optional[Any]:


  '''simple docstring'''






  lowercase__  :   Dict    		=      NestedDataStructure(lowercase_   ).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	(      lowercase_ , lowercase_   )		->	str:


  '''simple docstring'''






  lowercase__  :   Dict    		=      NestedDataStructure(lowercase_   ).flatten()
  assert output == expected_output

def 		UpperCamelCase	(      )		->	List[str]:


  '''simple docstring'''






  lowercase__  :   int    		=      A(x=1 , y="""foobar"""   )
  lowercase__  :   List[str]    		=      {"""x""": 1, """y""": """foobar"""}
  assert asdict(lowercase_   ) == expected_output

  lowercase__  :   Union[str, Any]    		=      {"""a""": {"""b""": A(x=10 , y="""foo"""   )}, """c""": [A(x=20 , y="""bar"""   )]}
  lowercase__  :   Union[str, Any]    		=      {"""a""": {"""b""": {"""x""": 10, """y""": """foo"""}}, """c""": [{"""x""": 20, """y""": """bar"""}]}
  assert asdict(lowercase_   ) == expected_output

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

def 		UpperCamelCase	(      lowercase_   )		->	int:


  '''simple docstring'''






  return text.split()

def 		UpperCamelCase	(      lowercase_   )		->	Union[str, Any]:


  '''simple docstring'''






  yield (time.time(), content)
  time.sleep(2   )
  yield (time.time(), content)



def 		UpperCamelCase	(      )		->	Tuple:


  '''simple docstring'''






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

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

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








from ..utils import DummyObject, requires_backends



class       _snake_case   (	metaclass=UpperCAmelCase_							):
 __lowerCAmelCase  :     int        =			['torch', 'transformers', 'onnx']


 def __init__(					self					,  *SCREAMING_SNAKE_CASE_					,  **SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




   requires_backends(self					,  ["""torch""", """transformers""", """onnx"""])


 @classmethod
 def      lowercase__     (					cls					,  *SCREAMING_SNAKE_CASE_					,  **SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




   requires_backends(cls					,  ["""torch""", """transformers""", """onnx"""])



 @classmethod
 def      lowercase__     (					cls					,  *SCREAMING_SNAKE_CASE_					,  **SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




   requires_backends(cls					,  ["""torch""", """transformers""", """onnx"""])



class       _snake_case   (	metaclass=UpperCAmelCase_							):
 __lowerCAmelCase  :     Any        =			['torch', 'transformers', 'onnx']


 def __init__(					self					,  *SCREAMING_SNAKE_CASE_					,  **SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




   requires_backends(self					,  ["""torch""", """transformers""", """onnx"""])


 @classmethod
 def      lowercase__     (					cls					,  *SCREAMING_SNAKE_CASE_					,  **SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




   requires_backends(cls					,  ["""torch""", """transformers""", """onnx"""])



 @classmethod
 def      lowercase__     (					cls					,  *SCREAMING_SNAKE_CASE_					,  **SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




   requires_backends(cls					,  ["""torch""", """transformers""", """onnx"""])



class       _snake_case   (	metaclass=UpperCAmelCase_							):
 __lowerCAmelCase  :     Union[str, Any]        =			['torch', 'transformers', 'onnx']


 def __init__(					self					,  *SCREAMING_SNAKE_CASE_					,  **SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




   requires_backends(self					,  ["""torch""", """transformers""", """onnx"""])


 @classmethod
 def      lowercase__     (					cls					,  *SCREAMING_SNAKE_CASE_					,  **SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




   requires_backends(cls					,  ["""torch""", """transformers""", """onnx"""])



 @classmethod
 def      lowercase__     (					cls					,  *SCREAMING_SNAKE_CASE_					,  **SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




   requires_backends(cls					,  ["""torch""", """transformers""", """onnx"""])



class       _snake_case   (	metaclass=UpperCAmelCase_							):
 __lowerCAmelCase  :     List[Any]        =			['torch', 'transformers', 'onnx']


 def __init__(					self					,  *SCREAMING_SNAKE_CASE_					,  **SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




   requires_backends(self					,  ["""torch""", """transformers""", """onnx"""])


 @classmethod
 def      lowercase__     (					cls					,  *SCREAMING_SNAKE_CASE_					,  **SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




   requires_backends(cls					,  ["""torch""", """transformers""", """onnx"""])



 @classmethod
 def      lowercase__     (					cls					,  *SCREAMING_SNAKE_CASE_					,  **SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




   requires_backends(cls					,  ["""torch""", """transformers""", """onnx"""])



class       _snake_case   (	metaclass=UpperCAmelCase_							):
 __lowerCAmelCase  :     Dict        =			['torch', 'transformers', 'onnx']


 def __init__(					self					,  *SCREAMING_SNAKE_CASE_					,  **SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




   requires_backends(self					,  ["""torch""", """transformers""", """onnx"""])


 @classmethod
 def      lowercase__     (					cls					,  *SCREAMING_SNAKE_CASE_					,  **SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




   requires_backends(cls					,  ["""torch""", """transformers""", """onnx"""])



 @classmethod
 def      lowercase__     (					cls					,  *SCREAMING_SNAKE_CASE_					,  **SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




   requires_backends(cls					,  ["""torch""", """transformers""", """onnx"""])



class       _snake_case   (	metaclass=UpperCAmelCase_							):
 __lowerCAmelCase  :     List[str]        =			['torch', 'transformers', 'onnx']


 def __init__(					self					,  *SCREAMING_SNAKE_CASE_					,  **SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




   requires_backends(self					,  ["""torch""", """transformers""", """onnx"""])


 @classmethod
 def      lowercase__     (					cls					,  *SCREAMING_SNAKE_CASE_					,  **SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




   requires_backends(cls					,  ["""torch""", """transformers""", """onnx"""])



 @classmethod
 def      lowercase__     (					cls					,  *SCREAMING_SNAKE_CASE_					,  **SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




   requires_backends(cls					,  ["""torch""", """transformers""", """onnx"""])






