In a python iterator select other iterator based on conditions
In python I have an iterator returning an infinite string of indices in a fixed range [0, N] called Sampler. Actually I have a list of those and all they do is return indices in the range [0, N_0], [N_0, N_1], ..., [N_{n-1}, N_n].
What I now want to do is first select one of these iterators based on the length of their range, so I have a weights list [N_0, N_1 - N_0, ...] and I select one of these with:
    iterator_idx = random.choices(range(len(weights)), weights=weights/weights.sum())[0]

Next, what I want to do is create an iterator which randomly selects one of the iterators and selects a batch of M samples.
class BatchSampler:
    def __init__(self, M):
        self.M = M
        self.weights = [weight_list]

        self.samplers = [list_of_iterators]
        ]
        self._batch_samplers = [
            self.batch_sampler(sampler) for sampler in self.samplers
        ]

    def batch_sampler(self, sampler):
        batch = []
        for batch_idx in sampler:
            batch.append(batch_idx)
            if len(batch) == self.M:
                yield batch

        if len(batch) &gt; 0:
            yield batch

    def __iter__(self):
        # First select one of the datasets.
        iterator_idx = random.choices(
            range(len(self.weights)), weights=self.weights / self.weights.sum()
        )[0]
        return self._batch_samplers[iterator_idx]

The issue with this is that iter() only seems to be called once, so only the first time iterator_idx is selected. Obviously this is wrong... What is the way around this?
This is a possible case when you would have multiple datasets in pytorch, but you want to sample only batches from one of the datasets.

Seems to me that you want to define your own container type.
I'll try to provide examples of a few standard ways to do so
(hopefully without missing too many details);
you should be able to reuse one of these simple examples,
into your own class.

Using just __getitem__ (support indexing &amp; looping):
object.__getitem__

Called to implement evaluation of self[key].


class MyContainer:
  def __init__(self, sequence):
    self.elements = sequence  # Just something to work with.
  
  def __getitem__(self, key):
    # If we're delegating to sequences like built-in list, 
    # invalid indices are handled automatically by them 
    # (throwing IndexError, as per the documentation).
    return self.elements[key]

t = (1, 2, 'a', 'b')
c = MyContainer(t)
elems = [e for e in c]
assert elems == [1, 2, 'a', 'b']
assert c[1:-1] == t[1:-1] == (2, 'a')



Using the iterator protocol:
object.__iter__

object.__iter__(self)
This method is called when an iterator is required for a container. This method should return a new iterator object that can iterate over all the objects in the container. For mappings, it should iterate over the keys of the container.
Iterator objects also need to implement this method; they are required to return themselves. For more information on iterator objects, see Iterator Types.

Iterator Types

container.__iter__()
Return an iterator object. The object is required to support the iterator protocol described below.


The iterator objects themselves are required to support the following two methods, which together form the iterator protocol:
iterator.__iter__()
Return the iterator object itself. This is required to allow both containers and iterators to be used with the for and in statements.
iterator.__next__()
Return the next item from the container. If there are no further items, raise the StopIteration exception.


Once an iterator's __next__() method raises StopIteration, it must continue to do so on subsequent calls.


class MyContainer:
  class Iter:
    def __init__(self, container):
      self.cont = container
      self.pos = 0
      self.len = len(container.elements)
    
    def __iter__(self): return self
    def __next__(self):
      if self.pos == self.len: raise StopIteration
      curElem = self.cont.elements[self.pos]
      self.pos += 1
      return curElem
  
  def __init__(self, sequence):
    self.elements = sequence  # Just something to work with.
  
  def __iter__(self):
    return MyContainer.Iter(self)

t = (1, 2, 'a', 'b')
c = MyContainer(t)
elems = [e for e in c]
assert elems == [1, 2, 'a', 'b']



Using a generator:
Generator Types

Python's generators provide a convenient way to implement the iterator protocol. If a container object's iter() method is implemented as a generator, it will automatically return an iterator object (technically, a generator object) supplying the iter() and next() methods.

generator

A function which returns a generator iterator. It looks like a normal function except that it contains yield expressions for producing a series of values usable in a for-loop or that can be retrieved one at a time with the next() function.
Usually refers to a generator function, but may refer to a generator iterator in some contexts.

generator iterator

An object created by a generator function.

6.2.9. Yield expressions

Using a yield expression in a function's body causes that function to be a generator


class MyContainer:
  def __init__(self, sequence):
    self.elements = sequence  # Just something to work with.
  
  def __iter__(self):
    for e in self.elements: yield e

t = (1, 2, 'a', 'b')
c = MyContainer(t)
elems = [e for e in c]
assert elems == [1, 2, 'a', 'b']


