



'''simple docstring'''



import multiprocessing
from typing import TYPE_CHECKING, Optional, Union

from .. import Dataset, Features, config
from ..formatting import query_table
from ..packaged_modules.sql.sql import Sql
from ..utils import logging
from .abc import AbstractDatasetInputStream


if TYPE_CHECKING:
		import sqlitea

		import sqlalchemy





class _UpperCamelCase		(     lowerCAmelCase__		):





				'''simple docstring'''





				def __init__(							self  :  Optional[Any]  ,       _lowerCAmelCase  :  List[Any]  ,       _lowerCAmelCase  :  Union[str, Any]  ,       _lowerCAmelCase  :  List[Any] = None  ,       _lowerCAmelCase  :  Optional[int] = None  ,       _lowerCAmelCase  :  Any = False  ,       **_lowerCAmelCase  :  Tuple  ,       ):







											'''simple docstring'''

											super().__init__(features=_SCREAMING_SNAKE_CASE  ,       cache_dir=_SCREAMING_SNAKE_CASE  ,       keep_in_memory=_SCREAMING_SNAKE_CASE  ,       **_SCREAMING_SNAKE_CASE)
											__lowercase       	=Sql(
											    cache_dir=_SCREAMING_SNAKE_CASE  ,       features=_SCREAMING_SNAKE_CASE  ,       sql=_SCREAMING_SNAKE_CASE  ,       con=_SCREAMING_SNAKE_CASE  ,       **_SCREAMING_SNAKE_CASE  ,       )

				def 				__lowerCamelCase     (							self  :  Union[str, Any]):







											'''simple docstring'''

											__lowercase       	=None
											__lowercase       	=None
											__lowercase       	=None
											__lowercase       	=None

											self.builder.download_and_prepare(
											    download_config=_SCREAMING_SNAKE_CASE  ,       download_mode=_SCREAMING_SNAKE_CASE  ,       verification_mode=_SCREAMING_SNAKE_CASE  ,       base_path=_SCREAMING_SNAKE_CASE  ,       )

											# Build dataset for splits
											__lowercase       	=self.builder.as_dataset(
											    split='train'  ,       verification_mode=_SCREAMING_SNAKE_CASE  ,       in_memory=self.keep_in_memory)
											return dataset





class _UpperCamelCase		:





				'''simple docstring'''





				def __init__(							self  :  Any  ,       _lowerCAmelCase  :  int  ,       _lowerCAmelCase  :  Dict  ,       _lowerCAmelCase  :  int  ,       _lowerCAmelCase  :  Union[str, Any] = None  ,       _lowerCAmelCase  :  Union[str, Any] = None  ,       **_lowerCAmelCase  :  List[Any]  ,       ):







											'''simple docstring'''

											if num_proc is not None and num_proc <= 0:
																		raise ValueError(f"""num_proc {num_proc} must be an integer > 0.""")

											__lowercase       	=dataset
											__lowercase       	=name
											__lowercase       	=con
											__lowercase       	=batch_size if batch_size else config.DEFAULT_MAX_BATCH_SIZE
											__lowercase       	=num_proc
											__lowercase       	=to_sql_kwargs

				def 				__lowerCamelCase     (							self  :  Any):







											'''simple docstring'''

											__lowercase       	=self.to_sql_kwargs.pop('sql'  ,       _SCREAMING_SNAKE_CASE)
											__lowercase       	=self.to_sql_kwargs.pop('con'  ,       _SCREAMING_SNAKE_CASE)
											__lowercase       	=self.to_sql_kwargs.pop('index'  ,       _SCREAMING_SNAKE_CASE)

											__lowercase       	=self._write(index=_SCREAMING_SNAKE_CASE  ,       **self.to_sql_kwargs)
											return written

				def 				__lowerCamelCase     (							self  :  Dict  ,       _lowerCAmelCase  :  Union[str, Any]):







											'''simple docstring'''

											__lowercase	,				__lowercase	,				__lowercase       	=args
											__lowercase       	={**to_sql_kwargs, 'if_exists': 'append'} if offset > 0 else to_sql_kwargs
											__lowercase       	=query_table(
											    table=self.dataset.data  ,       key=slice(_SCREAMING_SNAKE_CASE  ,       offset + self.batch_size)  ,       indices=self.dataset._indices  ,       )
											__lowercase       	=batch.to_pandas()
											__lowercase       	=df.to_sql(self.name  ,       self.con  ,       index=_SCREAMING_SNAKE_CASE  ,       **_SCREAMING_SNAKE_CASE)
											return num_rows or len(_SCREAMING_SNAKE_CASE)

				def 				__lowerCamelCase     (							self  :  List[Any]  ,       _lowerCAmelCase  :  Tuple  ,       **_lowerCAmelCase  :  Optional[int]):







											'''simple docstring'''

											__lowercase       	=0

											if self.num_proc is None or self.num_proc == 1:
																		for offset in logging.tqdm(
																		    range(0  ,       len(self.dataset)  ,       self.batch_size)  ,       unit='ba'  ,       disable=not logging.is_progress_bar_enabled()  ,       desc='Creating SQL from Arrow format'  ,       ):
																									written += self._batch_sql((offset, index, to_sql_kwargs))
											else:
																		__lowercase	,				__lowercase       	=len(self.dataset), self.batch_size
																		with multiprocessing.Pool(self.num_proc) as pool:
																									for num_rows in logging.tqdm(
																									    pool.imap(
																									        self._batch_sql  ,       [(offset, index, to_sql_kwargs) for offset in range(0  ,       _SCREAMING_SNAKE_CASE  ,       _SCREAMING_SNAKE_CASE)]  ,       )  ,       total=(num_rows // batch_size) + 1 if num_rows % batch_size else num_rows // batch_size  ,       unit='ba'  ,       disable=not logging.is_progress_bar_enabled()  ,       desc='Creating SQL from Arrow format'  ,       ):
																																written += num_rows

											return written








'''simple docstring'''



import os

import pytest
import yaml

from datasets.features.features import Features, Value
from datasets.info import DatasetInfo, DatasetInfosDict



@pytest.mark.parametrize(
    'files'  ,   [
        ['full:README.md', 'dataset_infos.json'],
        ['empty:README.md', 'dataset_infos.json'],
        ['dataset_infos.json'],
        ['full:README.md'],
    ]  ,   )
def   _A       (   _lowerCAmelCase  ,   _lowerCAmelCase					):



							"""simple docstring"""



							__lowercase       	=tmp_path_factory.mktemp('dset_infos_dir'					)
							if "full:README.md" in files:
														with open(dataset_infos_dir / 'README.md'  ,   'w'					) as f:
																					f.write('---\ndataset_info:\n  dataset_size: 42\n---'					)
							if "empty:README.md" in files:
														with open(dataset_infos_dir / 'README.md'  ,   'w'					) as f:
																					f.write(''					)
    # we want to support dataset_infos.json for backward compatibility
							if "dataset_infos.json" in files:
														with open(dataset_infos_dir / 'dataset_infos.json'  ,   'w'					) as f:
																					f.write('{"default": {"dataset_size": 42}}'					)
							__lowercase       	=DatasetInfosDict.from_directory(_lowerCAmelCase					)
							assert dataset_infos
							assert dataset_infos["default"].dataset_size == 42



@pytest.mark.parametrize(
    'dataset_info'  ,   [
        DatasetInfo(),
        DatasetInfo(
            description='foo'  ,   features=Features({'a': Value('int32'					)}					)  ,   builder_name='builder'  ,   config_name='config'  ,   version='1.0.0'  ,   splits=[{'name': 'train'}]  ,   download_size=42  ,   ),
    ]  ,   )
def   _A       (   _lowerCAmelCase  ,   _lowerCAmelCase					):



							"""simple docstring"""



							__lowercase       	=str(_lowerCAmelCase					)
							dataset_info.write_to_directory(_lowerCAmelCase					)
							__lowercase       	=DatasetInfo.from_directory(_lowerCAmelCase					)
							assert dataset_info == reloaded
							assert os.path.exists(os.path.join(_lowerCAmelCase  ,   'dataset_info.json'					)					)



def   _A       (   ):



							"""simple docstring"""



							__lowercase       	=DatasetInfo(
							    description='foo'  ,   citation='bar'  ,   homepage='https://foo.bar'  ,   license='CC0'  ,   features=Features({'a': Value('int32'					)}					)  ,   post_processed={}  ,   supervised_keys=()  ,   task_templates=[]  ,   builder_name='builder'  ,   config_name='config'  ,   version='1.0.0'  ,   splits=[{'name': 'train', 'num_examples': 42}]  ,   download_checksums={}  ,   download_size=1_337  ,   post_processing_size=442  ,   dataset_size=1_234  ,   size_in_bytes=1_337 + 442 + 1_234  ,   )
							__lowercase       	=dataset_info._to_yaml_dict()
							assert sorted(_lowerCAmelCase					) == sorted(DatasetInfo._INCLUDED_INFO_IN_YAML					)
							for key in DatasetInfo._INCLUDED_INFO_IN_YAML:
														assert key in dataset_info_yaml_dict
														assert isinstance(dataset_info_yaml_dict[key]  ,   (list, dict, int, str)					)
							__lowercase       	=yaml.safe_dump(_lowerCAmelCase					)
							__lowercase       	=yaml.safe_load(_lowerCAmelCase					)
							assert dataset_info_yaml_dict == reloaded



def   _A       (   ):



							"""simple docstring"""



							__lowercase       	=DatasetInfo()
							__lowercase       	=dataset_info._to_yaml_dict()
							assert dataset_info_yaml_dict == {}



@pytest.mark.parametrize(
    'dataset_infos_dict'  ,   [
        DatasetInfosDict(),
        DatasetInfosDict({'default': DatasetInfo()}					),
        DatasetInfosDict({'my_config_name': DatasetInfo()}					),
        DatasetInfosDict(
            {
                'default': DatasetInfo(
                    description='foo'  ,   features=Features({'a': Value('int32'					)}					)  ,   builder_name='builder'  ,   config_name='config'  ,   version='1.0.0'  ,   splits=[{'name': 'train'}]  ,   download_size=42  ,   )
            }					),
        DatasetInfosDict(
            {
                'v1': DatasetInfo(dataset_size=42					),
                'v2': DatasetInfo(dataset_size=1_337					),
            }					),
    ]  ,   )
def   _A       (   _lowerCAmelCase  ,   _lowerCAmelCase					):



							"""simple docstring"""



							__lowercase       	=str(_lowerCAmelCase					)
							dataset_infos_dict.write_to_directory(_lowerCAmelCase					)
							__lowercase       	=DatasetInfosDict.from_directory(_lowerCAmelCase					)

							# the config_name of the dataset_infos_dict take over the attribute
							for config_name, dataset_info in dataset_infos_dict.items():
														__lowercase       	=config_name
														# the yaml representation doesn't include fields like description or citation
														# so we just test that we can recover what we can from the yaml
														__lowercase       	=DatasetInfo._from_yaml_dict(dataset_info._to_yaml_dict()					)
							assert dataset_infos_dict == reloaded

							if dataset_infos_dict:
														assert os.path.exists(os.path.join(_lowerCAmelCase  ,   'README.md'					)					)



