<commit_before># pylint: disable=missing-docstring, redefined-outer-name

import pytest
import numpy as np

from batchflow import Dataset, Batch, ImagesBatch, DatasetIndex, Pipeline


@pytest.fixture
def dataset():
    index = DatasetIndex(100)
    return Dataset(index, Batch)


class TestDataset:
    def test_from_dataset(self, dataset):
        new_index = DatasetIndex(25)
        new_ds = dataset.from_dataset(dataset, new_index)
        assert isinstance(new_ds, dataset.__class__)
        assert new_ds.batch_class == dataset.batch_class
        assert (new_ds.index.index == new_index.index).all()

    def test_from_dataset_new_batch_class(self, dataset):
        new_index = DatasetIndex(25)
        new_ds = Dataset.from_dataset(dataset, new_index, batch_class=ImagesBatch)
        assert isinstance(new_ds, dataset.__class__)
        assert new_ds.batch_class == ImagesBatch

    def test_build_index(self):
        new_index = Dataset.build_index(25)
        assert isinstance(new_index, DatasetIndex)

    def test_create_subset(self, dataset):
        new_index = DatasetIndex(25)
        new_ds = dataset.create_subset(new_index)
        assert isinstance(new_ds, dataset.__class__)
        assert np.isin(new_ds.indices, dataset.indices).all()

    def test_create_subset_wrong_index(self, dataset):
        wrong_index = DatasetIndex(np.arange(200, 225))
        with pytest.raises(IndexError) as error:
            dataset.create_subset(wrong_index)
            assert 'IndexError' in str(error)

    def test_create_batch(self, dataset):
        target_index = DatasetIndex(5)
        new_batch = dataset.create_batch(target_index)
        assert isinstance(new_batch, dataset.batch_class)
        assert len(new_batch.indices) == len(target_index.indices)

    def test_pipeline(self, dataset):
        pipeline_config = {}
        new_pipeline = dataset.pipeline(pipeline_config)
        assert isinstance(new_pipeline, Pipeline)

    def test_rshift(self, dataset):
        pipeline_config = {}
        new_pipeline = dataset.pipeline(pipeline_config)
        train_pipeline = (new_pipeline << dataset)
        assert isinstance(train_pipeline, Pipeline)

    def test_split(self, dataset):
        assert dataset.train is None
        dataset.split((0.8, 0.2))
        assert dataset.train is not None
        assert dataset.test is not None
        assert dataset.validation is None
<commit_msg>Add more tests for Dataset-class split method<commit_after># pylint: disable=missing-docstring, redefined-outer-name

import pytest
import numpy as np

from batchflow import Dataset, Batch, ImagesBatch, DatasetIndex, Pipeline


@pytest.fixture
def dataset():
    index = DatasetIndex(100)
    return Dataset(index, Batch)


class TestDataset:
    def test_from_dataset(self, dataset):
        new_index = DatasetIndex(25)
        new_ds = dataset.from_dataset(dataset, new_index)
        assert isinstance(new_ds, dataset.__class__)
        assert new_ds.batch_class == dataset.batch_class
        assert (new_ds.index.index == new_index.index).all()

    def test_from_dataset_new_batch_class(self, dataset):
        new_index = DatasetIndex(25)
        new_ds = Dataset.from_dataset(dataset, new_index, batch_class=ImagesBatch)
        assert isinstance(new_ds, dataset.__class__)
        assert new_ds.batch_class == ImagesBatch

    def test_build_index(self):
        new_index = Dataset.build_index(25)
        assert isinstance(new_index, DatasetIndex)

    def test_create_subset(self, dataset):
        new_index = DatasetIndex(25)
        new_ds = dataset.create_subset(new_index)
        assert isinstance(new_ds, dataset.__class__)
        assert np.isin(new_ds.indices, dataset.indices).all()

    def test_create_subset_wrong_index(self, dataset):
        wrong_index = DatasetIndex(np.arange(200, 225))
        with pytest.raises(IndexError) as error:
            dataset.create_subset(wrong_index)
            assert 'IndexError' in str(error)

    def test_create_batch(self, dataset):
        target_index = DatasetIndex(5)
        new_batch = dataset.create_batch(target_index)
        assert isinstance(new_batch, dataset.batch_class)
        assert len(new_batch.indices) == len(target_index.indices)

    def test_pipeline(self, dataset):
        pipeline_config = {}
        new_pipeline = dataset.pipeline(pipeline_config)
        assert isinstance(new_pipeline, Pipeline)

    def test_rshift(self, dataset):
        pipeline_config = {}
        new_pipeline = dataset.pipeline(pipeline_config)
        train_pipeline = (new_pipeline << dataset)
        assert isinstance(train_pipeline, Pipeline)

    def test_split(self, dataset):
        assert dataset.train is None
        train_part, test_part = 0.8, 0.2
        dataset.split([train_part, test_part])
        assert dataset.train is not None
        assert dataset.test is not None
        assert dataset.validation is None
        assert dataset.train.indices.shape[0]/dataset.indices.shape[0] == train_part
        assert dataset.test.indices.shape[0]/dataset.indices.shape[0] == test_part
        train_part, test_part, validation_part = 0.7, 0.2, 0.1
        dataset.split([train_part, test_part])
        assert dataset.validation is not None
        assert dataset.train.indices.shape[0] / dataset.indices.shape[0] == train_part
        assert dataset.test.indices.shape[0] / dataset.indices.shape[0] == test_part
        assert dataset.validation.indices.shape[0] / dataset.indices.shape[0] == validation_part

{'fraction_non_alphanumeric': 0.08857091106683572, 'fraction_numerical': 0.010867596449918493, 'mean_word_length': 3.320813771517997, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '903539', 'n_tokens_mistral': 1649, 'n_tokens_neox': 1626, 'n_words': 345}