如何在python中实现自定义数据的next_batch()函数

class DataIterator:
    def __init__(self, data):
        self.data = data
        self.start_index = 0

    def next_batch(self, batch_size):
        end_index = self.start_index + batch_size
        batch_data = self.data[self.start_index:end_index]

        if end_index >= len(self.data):
            self.start_index = 0
        else:
            self.start_index = end_index

        return batch_data

data = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
iterator = DataIterator(data)

batch = iterator.next_batch(3)
print(batch)  # 输出：[1, 2, 3]

batch = iterator.next_batch(4)
print(batch)  # 输出：[4, 5, 6, 7]

batch = iterator.next_batch(5)
print(batch)  # 输出：[8, 9, 10]