How would an iterator equivalent to itertools.count be implemented in Python without using a generator function?

Question

With a generator function, this is how itertools.count can be implemented (from documentation):

def count(start=0, step=1):
    # count(10) --> 10 11 12 13 14 ...
    # count(2.5, 0.5) -> 2.5 3.0 3.5 ...
    n = start
    while True:
        yield n
        n += step

I am trying to find how a similar iterator could be implemented without a generator function.

class Count:
    def __init__(self, start=0, step=1):
        self.c = start
        self.step = step

    def __iter__(self):
        return self

    def __next__(self):
        n = self.c
        self.c += self.step
        return n

Is this the implementation?

For sure, it does not make any practical sense, the point is just to get an idea of how it is achieved with generators.

Looks ok to me.

Michael Butscher
– Michael Butscher

2019-02-01 21:24:03 +00:00
Commented Feb 1, 2019 at 21:24 — Michael Butscher
– Michael Butscher, Commented Feb 1, 2019 at 21:24

Victor Ruiz · Accepted Answer · 2019-07-03 09:56:22Z

1

You could use collections.abc.Iterator to implement count with less code

from collections.abc import Iterator

class count(Iterator):
    def __init__(self, start=0, step=1):
        self.c, self.step = start-step, step

    def __next__(self):
        self.c += self.step
        return self.c

edited Jul 3, 2019 at 9:56

answered Jul 3, 2019 at 9:32

Victor Ruiz

1,21210 silver badges11 bronze badges

Sign up to request clarification or add additional context in comments.

Comments

Laurent LAPORTE · Accepted Answer · 2019-02-01 22:33:49Z

One of the best practices is to write a small unit test. You can use the unittest framework for that.

Here is an example with several tests which use your class in a loop and check the values of your iterator:

import unittest
import random


class TestCount(unittest.TestCase):
    loops = 20

    def test_default(self):
        c = Count()
        for i in range(self.loops):
            self.assertEqual(i, next(c))

    def test_start(self):
        start = random.randint(-10, 10)
        c = Count(start=start)
        for i in range(start, start + self.loops):
            self.assertEqual(i, next(c))

    def test_step_pos(self):
        step = random.randint(1, 5)
        c = Count(step=step)
        for i in range(0, self.loops, step):
            self.assertEqual(i, next(c))

    def test_step_neg(self):
        step = random.randint(-5, -1)
        c = Count(step=step)
        for i in range(0, -self.loops, step):
            self.assertEqual(i, next(c))

Collectives™ on Stack Overflow

How would an iterator equivalent to itertools.count be implemented in Python without using a generator function?

2 Answers 2

Comments

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

Comments

Comments

Your Answer

Sign up or log in

Post as a guest

Related