


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





def __snake_case   ( UpperCAmelCase_		:		Any						):  # picklable for multiprocessing
 return i + 1





@dataclass
class 	snake_case       :





  """simple docstring"""




  _lowerCamelCase						       =    42
  _lowerCamelCase						       =    42





class 	snake_case       (       lowercase     ):





  """simple docstring"""




  def     snake_case						(       self	):







   """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(lowercase__     ,    lowercase__	)     ,    lowercase__	)
   self.assertEqual(map_nested(lowercase__     ,    lowercase__	)     ,    lowercase__	)
   self.assertEqual(map_nested(lowercase__     ,    lowercase__	)     ,    lowercase__	)
   self.assertEqual(map_nested(lowercase__     ,    lowercase__	)     ,    lowercase__	)
   self.assertEqual(map_nested(lowercase__     ,    lowercase__	)     ,    lowercase__	)
   self.assertEqual(map_nested(lowercase__     ,    lowercase__	)     ,    lowercase__	)
   self.assertEqual(map_nested(lowercase__     ,    lowercase__	)     ,    lowercase__	)
   self.assertEqual(map_nested(lowercase__     ,    lowercase__	)     ,    lowercase__	)

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

   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(lowercase__	),
       "b": np.zeros(3	).astype(lowercase__	),
       "c": np.ones(2	).astype(lowercase__	),
   }
   self.assertEqual(map_nested(lowercase__     ,    lowercase__     ,    map_numpy=lowercase__	)     ,    lowercase__	)
   self.assertEqual(
       {k: v.tolist() for k, v in map_nested(lowercase__     ,    lowercase__     ,    map_numpy=lowercase__	).items()}     ,    {k: v.tolist() for k, v in expected_map_nested_sna_int.items()}     ,    )
   self.assertEqual(map_nested(lowercase__     ,    lowercase__     ,    map_numpy=lowercase__     ,    num_proc=lowercase__	)     ,    lowercase__	)
   self.assertEqual(
       {k: v.tolist() for k, v in map_nested(lowercase__     ,    lowercase__     ,    map_numpy=lowercase__     ,    num_proc=lowercase__	).items()}     ,    {k: v.tolist() for k, v in expected_map_nested_sna_int.items()}     ,    )
   with self.assertRaises(lowercase__	):  # can't pickle a local lambda
    map_nested(lambda UpperCamelCase	: x + 1     ,    lowercase__     ,    num_proc=lowercase__	)






  def     snake_case						(       self	):







   """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(lowercase__     ,    lowercase__     ,    lowercase__	)	)     ,    lowercase__	)







  def     snake_case						(       self	):







   """simple docstring"""





   class 	snake_case       :





     """simple docstring"""




     _lowerCamelCase						       =    "bar"



   lowerCamelCase_						=    Foo()
   self.assertEqual(foo.my_attr     ,    "bar"	)
   with temporary_assignment(lowercase__     ,    "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 __snake_case   ( UpperCAmelCase_		:		int   ,							UpperCAmelCase_		:		List[Any]   ,							UpperCAmelCase_		:		Union[str, Any]						):
 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(__A						)}
  lowerCamelCase_						=    map_nested(lambda UpperCAmelCase_						: 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 	snake_case       (       lowercase     ):





  """simple docstring"""




  @require_tf
  def     snake_case						(       self	):







   """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(lowercase__	).numpy()

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

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






  @require_torch
  def     snake_case						(       self	):







   """simple docstring"""
   import torch

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

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

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







  def     snake_case						(       self	):







   """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(lowercase__     ,    lowercase__	)
   self.assertGreater(np.abs(outa - outa	).sum()     ,    0	)





@pytest.mark.parametrize("input_data"   ,							[{}]						)
def __snake_case   ( UpperCAmelCase_		:		List[Any]						):
 lowerCamelCase_						=    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 __snake_case   ( UpperCAmelCase_		:		List[Any]   ,							UpperCAmelCase_		:		Tuple						):
 lowerCamelCase_						=    NestedDataStructure(__A						).flatten()
 assert output == expected_output





def __snake_case   ( ):
 lowerCamelCase_						=    A(x=1   ,							y="foobar"						)
 lowerCamelCase_						=    {"x": 1, "y": "foobar"}
 assert asdict(__A						) == 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(__A						) == expected_output

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





def __snake_case   ( UpperCAmelCase_		:		str						):
 return text.split()





def __snake_case   ( UpperCAmelCase_		:		int						):
 yield (time.time(), content)
 time.sleep(2						)
 yield (time.time(), content)







def __snake_case   ( ):
 with Pool(2						) as pool:
  lowerCamelCase_						=    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:
  lowerCamelCase_						=    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:
  lowerCamelCase_						=    []
  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






import inspect
import re
from hashlib import shaaaa
from typing import Dict, List

from .arrow import arrow
from .audiofolder import audiofolder
from .csv import csv
from .imagefolder import imagefolder
from .json import json
from .pandas import pandas
from .parquet import parquet
from .sql import sql  # noqa F401
from .text import text





def      A__					(			__A  :						List[str]			)					->str:
      __A    =[]
      for line in lines:
            __A    =re.sub(r'''#.*'''     ,				''''''     ,				__A			)  # remove comments
            if line:
                  filtered_lines.append(__A			)
      __A    ='''\n'''.join(__A			)

      # Make a hash from all this code
      __A    =full_str.encode('''utf-8'''			)
      return shaaaa(__A			).hexdigest()


# get importable module names and hash for caching
_lowerCamelCase      :					Dict					    =      {
    '''csv''': (csv.__name__, _hash_python_lines(inspect.getsource(csv).splitlines())),
    '''json''': (json.__name__, _hash_python_lines(inspect.getsource(json).splitlines())),
    '''pandas''': (pandas.__name__, _hash_python_lines(inspect.getsource(pandas).splitlines())),
    '''parquet''': (parquet.__name__, _hash_python_lines(inspect.getsource(parquet).splitlines())),
    '''arrow''': (arrow.__name__, _hash_python_lines(inspect.getsource(arrow).splitlines())),
    '''text''': (text.__name__, _hash_python_lines(inspect.getsource(text).splitlines())),
    '''imagefolder''': (imagefolder.__name__, _hash_python_lines(inspect.getsource(imagefolder).splitlines())),
    '''audiofolder''': (audiofolder.__name__, _hash_python_lines(inspect.getsource(audiofolder).splitlines())),
}

# Used to infer the module to use based on the data files extensions
_lowerCamelCase      :					int					    =      {
    '''.csv''': ('''csv''', {}),
    '''.tsv''': ('''csv''', {'''sep''': '''\t'''}),
    '''.json''': ('''json''', {}),
    '''.jsonl''': ('''json''', {}),
    '''.parquet''': ('''parquet''', {}),
    '''.arrow''': ('''arrow''', {}),
    '''.txt''': ('''text''', {}),
}
_EXTENSION_TO_MODULE.update({ext: ('''imagefolder''', {}) for ext in imagefolder.ImageFolder.EXTENSIONS})
_EXTENSION_TO_MODULE.update({ext.upper(): ('''imagefolder''', {}) for ext in imagefolder.ImageFolder.EXTENSIONS})
_EXTENSION_TO_MODULE.update({ext: ('''audiofolder''', {}) for ext in audiofolder.AudioFolder.EXTENSIONS})
_EXTENSION_TO_MODULE.update({ext.upper(): ('''audiofolder''', {}) for ext in audiofolder.AudioFolder.EXTENSIONS})
_lowerCamelCase      :					List[str]					    =      {'''imagefolder''', '''audiofolder'''}

# Used to filter data files based on extensions given a module name
_lowerCamelCase      :					Dict[str, List[str]]					    =      {}
for _ext, (_module, _) in _EXTENSION_TO_MODULE.items():
   _MODULE_TO_EXTENSIONS.setdefault(_module, []).append(_ext)

_MODULE_TO_EXTENSIONS["imagefolder"].append('''.zip''')


_MODULE_TO_EXTENSIONS["audiofolder"].append('''.zip''')


