





"""simple docstring"""





import os
import tempfile
from pathlib import Path
from unittest import TestCase

import pyarrow as pa
import pytest

from datasets.arrow_dataset import Dataset
from datasets.arrow_reader import ArrowReader, BaseReader, FileInstructions, ReadInstruction, make_file_instructions
from datasets.info import DatasetInfo
from datasets.splits import NamedSplit, Split, SplitDict, SplitInfo

from .utils import assert_arrow_memory_doesnt_increase, assert_arrow_memory_increases

class 		lowercase ( _a ):





      def    _snake_case      (	self				,						lowercase				,						lowercase=False			)		->					int:
            lowerCAmelCase  						=	(
                filename_skip_take["filename"],
                filename_skip_take["skip"] if "skip" in filename_skip_take else None,
                filename_skip_take["take"] if "take" in filename_skip_take else None,
            )
            open(os.path.join(SCREAMING_SNAKE_CASE__			)				,						"""wb"""			).close()
            lowerCAmelCase  						=	pa.Table.from_pydict({"""filename""": [Path(SCREAMING_SNAKE_CASE__			).name] * 100}			)
            if take == -1:
                  lowerCAmelCase  						=	len(SCREAMING_SNAKE_CASE__			) - skip
            if skip is not None and take is not None:
                  lowerCAmelCase  						=	pa_table.slice(SCREAMING_SNAKE_CASE__				,						SCREAMING_SNAKE_CASE__			)
            return pa_table







class 		lowercase ( _a ):





      def    _snake_case      (	self			)		->					Dict:
            lowerCAmelCase  						=	"my_name"
            lowerCAmelCase  						=	SplitInfo(name="""train"""				,						num_examples=100			)
            lowerCAmelCase  						=	SplitInfo(name="""test"""				,						num_examples=100			)
            lowerCAmelCase  						=	[train_info, test_info]
            lowerCAmelCase  						=	SplitDict()
            split_dict.add(SCREAMING_SNAKE_CASE__			)
            split_dict.add(SCREAMING_SNAKE_CASE__			)
            lowerCAmelCase  						=	DatasetInfo(splits=SCREAMING_SNAKE_CASE__			)

            with tempfile.TemporaryDirectory() as tmp_dir:
                  lowerCAmelCase  						=	ReaderTest(SCREAMING_SNAKE_CASE__				,						SCREAMING_SNAKE_CASE__			)

                  lowerCAmelCase  						=	"test[:33%]"
                  lowerCAmelCase  						=	Dataset(**reader.read(SCREAMING_SNAKE_CASE__				,						SCREAMING_SNAKE_CASE__				,						SCREAMING_SNAKE_CASE__			)			)
                  self.assertEqual(dset["""filename"""][0]				,						f'{name}-test'			)
                  self.assertEqual(dset.num_rows				,						33			)
                  self.assertEqual(dset.num_columns				,						1			)

                  lowerCAmelCase  						=	["train", "test[:33%]"]
                  lowerCAmelCase  						=	[Split.TRAIN, ReadInstruction.from_spec("""test[:33%]"""			)]
                  for instructions in [instructionsa, instructionsa]:
                        lowerCAmelCase  						=	[reader.read(SCREAMING_SNAKE_CASE__				,						SCREAMING_SNAKE_CASE__				,						SCREAMING_SNAKE_CASE__			) for instr in instructions]
                        lowerCAmelCase  						=	(Dataset(**SCREAMING_SNAKE_CASE__			) for dataset_kwargs in datasets_kwargs)
                        self.assertEqual(train_dset["""filename"""][0]				,						f'{name}-train'			)
                        self.assertEqual(train_dset.num_rows				,						100			)
                        self.assertEqual(train_dset.num_columns				,						1			)
                        self.assertIsInstance(train_dset.split				,						SCREAMING_SNAKE_CASE__			)
                        self.assertEqual(str(train_dset.split			)				,						"""train"""			)
                        self.assertEqual(test_dset["""filename"""][0]				,						f'{name}-test'			)
                        self.assertEqual(test_dset.num_rows				,						33			)
                        self.assertEqual(test_dset.num_columns				,						1			)
                        self.assertIsInstance(test_dset.split				,						SCREAMING_SNAKE_CASE__			)
                        self.assertEqual(str(test_dset.split			)				,						"""test[:33%]"""			)
                        del train_dset, test_dset





      def    _snake_case      (	self			)		->					Any:
            lowerCAmelCase  						=	"my_name"
            lowerCAmelCase  						=	SplitInfo(name="""train"""				,						num_examples=1_000				,						shard_lengths=[100] * 10			)
            lowerCAmelCase  						=	[train_info]
            lowerCAmelCase  						=	SplitDict()
            split_dict.add(SCREAMING_SNAKE_CASE__			)
            lowerCAmelCase  						=	DatasetInfo(splits=SCREAMING_SNAKE_CASE__			)

            with tempfile.TemporaryDirectory() as tmp_dir:
                  lowerCAmelCase  						=	ReaderTest(SCREAMING_SNAKE_CASE__				,						SCREAMING_SNAKE_CASE__			)

                  lowerCAmelCase  						=	"train[:33%]"
                  lowerCAmelCase  						=	Dataset(**reader.read(SCREAMING_SNAKE_CASE__				,						SCREAMING_SNAKE_CASE__				,						SCREAMING_SNAKE_CASE__			)			)
                  self.assertEqual(dset["""filename"""][0]				,						f'{name}-train-00000-of-00010'			)
                  self.assertEqual(dset["""filename"""][-1]				,						f'{name}-train-00003-of-00010'			)
                  self.assertEqual(dset.num_rows				,						330			)
                  self.assertEqual(dset.num_columns				,						1			)





      def    _snake_case      (	self			)		->					Dict:
            lowerCAmelCase  						=	SplitInfo(name="""train"""				,						num_examples=100			)
            lowerCAmelCase  						=	SplitInfo(name="""test"""				,						num_examples=100			)
            lowerCAmelCase  						=	SplitDict()
            split_dict.add(SCREAMING_SNAKE_CASE__			)
            split_dict.add(SCREAMING_SNAKE_CASE__			)
            lowerCAmelCase  						=	DatasetInfo(splits=SCREAMING_SNAKE_CASE__			)

            with tempfile.TemporaryDirectory() as tmp_dir:
                  lowerCAmelCase  						=	ReaderTest(SCREAMING_SNAKE_CASE__				,						SCREAMING_SNAKE_CASE__			)

                  lowerCAmelCase  						=	[
                      {"filename": os.path.join(SCREAMING_SNAKE_CASE__				,						"""train"""			)},
                      {"filename": os.path.join(SCREAMING_SNAKE_CASE__				,						"""test"""			), "skip": 10, "take": 10},
                  ]
                  lowerCAmelCase  						=	Dataset(**reader.read_files(SCREAMING_SNAKE_CASE__				,						original_instructions="""train+test[10:20]"""			)			)
                  self.assertEqual(dset.num_rows				,						110			)
                  self.assertEqual(dset.num_columns				,						1			)
                  del dset







@pytest.mark.parametrize("""in_memory"""							,						[False, True]      )
def       UpperCAmelCase__							(  SCREAMING_SNAKE_CASE					:						Any							,						SCREAMING_SNAKE_CASE					:						int							,						SCREAMING_SNAKE_CASE					:						Optional[int]      ) ->    Optional[int]:







      '''simple docstring'''




      lowerCAmelCase  						=	arrow_file
      with assert_arrow_memory_increases() if in_memory else assert_arrow_memory_doesnt_increase():
            lowerCAmelCase  						=	ArrowReader.read_table(SCREAMING_SNAKE_CASE							,						in_memory=SCREAMING_SNAKE_CASE      )
      assert table.shape == dataset.data.shape
      assert set(table.column_names      ) == set(dataset.data.column_names      )
      assert dict(table.to_pydict()      ) == dict(dataset.data.to_pydict()      )  # to_pydict returns OrderedDict


@pytest.mark.parametrize("""in_memory"""							,						[False, True]      )
def       UpperCAmelCase__							(  SCREAMING_SNAKE_CASE					:						Optional[Any]							,						SCREAMING_SNAKE_CASE					:						Any							,						SCREAMING_SNAKE_CASE					:						Union[str, Any]      ) ->    List[str]:







      '''simple docstring'''




      lowerCAmelCase  						=	arrow_file
      lowerCAmelCase  						=	ArrowReader(""""""							,						SCREAMING_SNAKE_CASE      )
      with assert_arrow_memory_increases() if in_memory else assert_arrow_memory_doesnt_increase():
            lowerCAmelCase  						=	reader.read_files([{"""filename""": filename}]							,						in_memory=SCREAMING_SNAKE_CASE      )
      assert dataset_kwargs.keys() == {"arrow_table", "info", "split"}
      lowerCAmelCase  						=	dataset_kwargs["arrow_table"]
      assert table.shape == dataset.data.shape
      assert set(table.column_names      ) == set(dataset.data.column_names      )
      assert dict(table.to_pydict()      ) == dict(dataset.data.to_pydict()      )  # to_pydict returns OrderedDict


def       UpperCAmelCase__							(  ) ->    int:







      '''simple docstring'''




      assert ReadInstruction("""train"""							,						to=10							,						unit="""abs"""      ).to_spec() == "train[:10]"
      assert ReadInstruction("""train"""							,						from_=-80							,						to=10							,						unit="""%"""      ).to_spec() == "train[-80%:10%]"

      lowerCAmelCase  						=	"train+test"
      assert ReadInstruction.from_spec(SCREAMING_SNAKE_CASE      ).to_spec() == spec_train_test

      lowerCAmelCase  						=	"train[2:10]"
      assert ReadInstruction.from_spec(SCREAMING_SNAKE_CASE      ).to_spec() == spec_train_abs

      lowerCAmelCase  						=	"train[15%:-20%]"
      assert ReadInstruction.from_spec(SCREAMING_SNAKE_CASE      ).to_spec() == spec_train_pct

      lowerCAmelCase  						=	"train[:10%](closest)"
      assert ReadInstruction.from_spec(SCREAMING_SNAKE_CASE      ).to_spec() == "train[:10%]"

      lowerCAmelCase  						=	"train[:10%](pct1_dropremainder)"
      assert ReadInstruction.from_spec(SCREAMING_SNAKE_CASE      ).to_spec() == spec_train_pct_rounding

      lowerCAmelCase  						=	"train[:10%](pct1_dropremainder)+test[-10%:](pct1_dropremainder)"
      assert ReadInstruction.from_spec(SCREAMING_SNAKE_CASE      ).to_spec() == spec_train_test_pct_rounding


def       UpperCAmelCase__							(  ) ->    List[str]:







      '''simple docstring'''




      lowerCAmelCase  						=	"dummy"
      lowerCAmelCase  						=	[SplitInfo(name="""train"""							,						num_examples=1_00      )]
      lowerCAmelCase  						=	"train[:33%]"
      lowerCAmelCase  						=	"arrow"
      lowerCAmelCase  						=	"prefix"

      lowerCAmelCase  						=	make_file_instructions(SCREAMING_SNAKE_CASE							,						SCREAMING_SNAKE_CASE							,						SCREAMING_SNAKE_CASE							,						SCREAMING_SNAKE_CASE							,						SCREAMING_SNAKE_CASE      )
      assert isinstance(SCREAMING_SNAKE_CASE							,						SCREAMING_SNAKE_CASE      )
      assert file_instructions.num_examples == 33
      assert file_instructions.file_instructions == [
          {"filename": os.path.join(SCREAMING_SNAKE_CASE							,						F'{name}-train.arrow'      ), "skip": 0, "take": 33}
      ]

      lowerCAmelCase  						=	[SplitInfo(name="""train"""							,						num_examples=1_00							,						shard_lengths=[10] * 10      )]
      lowerCAmelCase  						=	make_file_instructions(SCREAMING_SNAKE_CASE							,						SCREAMING_SNAKE_CASE							,						SCREAMING_SNAKE_CASE							,						SCREAMING_SNAKE_CASE							,						SCREAMING_SNAKE_CASE      )
      assert isinstance(SCREAMING_SNAKE_CASE							,						SCREAMING_SNAKE_CASE      )
      assert file_instructions.num_examples == 33
      assert file_instructions.file_instructions == [
          {"filename": os.path.join(SCREAMING_SNAKE_CASE							,						F'{name}-train-00000-of-00010.arrow'      ), "skip": 0, "take": -1},
          {"filename": os.path.join(SCREAMING_SNAKE_CASE							,						F'{name}-train-00001-of-00010.arrow'      ), "skip": 0, "take": -1},
          {"filename": os.path.join(SCREAMING_SNAKE_CASE							,						F'{name}-train-00002-of-00010.arrow'      ), "skip": 0, "take": -1},
          {"filename": os.path.join(SCREAMING_SNAKE_CASE							,						F'{name}-train-00003-of-00010.arrow'      ), "skip": 0, "take": 3},
      ]






@pytest.mark.parametrize("""name, expected_exception"""							,						[(None, TypeError), ("""""", ValueError)]      )
def       UpperCAmelCase__							(  SCREAMING_SNAKE_CASE					:						Union[str, Any]							,						SCREAMING_SNAKE_CASE					:						Any      ) ->    List[str]:







      '''simple docstring'''




      lowerCAmelCase  						=	[SplitInfo(name="""train"""							,						num_examples=1_00      )]
      lowerCAmelCase  						=	"train"
      lowerCAmelCase  						=	"arrow"
      lowerCAmelCase  						=	"prefix_path"
      with pytest.raises(SCREAMING_SNAKE_CASE      ):
            lowerCAmelCase  						=	make_file_instructions(SCREAMING_SNAKE_CASE							,						SCREAMING_SNAKE_CASE							,						SCREAMING_SNAKE_CASE							,						SCREAMING_SNAKE_CASE							,						SCREAMING_SNAKE_CASE      )













"""simple docstring"""





from scipy.constants import g


def       UpperCAmelCase__							(  SCREAMING_SNAKE_CASE					:						float							,						SCREAMING_SNAKE_CASE					:						float      ):







      '''simple docstring'''




      if mass < 0:
            # handling of negative values of mass
            raise ValueError("""The mass of a body cannot be negative"""      )
      if height < 0:
            # handling of negative values of height
            raise ValueError("""The height above the ground cannot be negative"""      )
      return mass * g * height


if __name__ == "__main__":
       from doctest import testmod

       testmod(name="potential_energy")






