Calculating the square numbers within a range (python)

Question

I want to be able to execute the following code:

for i in Squares(5, 50):
      print(i)

Now this is very easy to implement using a loop, however I want to use an iterator.

So I have defined the following class:

import math

class Squares(object):

    def __init__(self, start, stop):
       self.start = start
       self.stop = stop

    def __iter__(self): 
        return self

    def __next__(self):
        start = self.start
        stop = self.stop
        squareroot = math.sqrt(start)

        if self.start > self.stop:
            raise StopIteration

        if squareroot == math.ceil(squareroot):
            start += 1

But at the moment this is returning None an infinite amount of times. This means the none must be because the StopIteration is being executed even when it shouldn't. I think my if squareroot == math.ceil(squareroot): condition is correct because I tested it separately, but I can't figure out what to change to get the output I want. Any help is appreciated.

EDIT: For a code such as:

for i in Squares(4, 16):
    print(i)

I would expect the output to be:

4
9
16

Answer 1

Try creating a generator function:

from math import sqrt, ceil

def Squares(start, stop):
    for i in range(start, stop+1):
        sqrti = sqrt(i)
        if sqrti == ceil(sqrti):
            yield i

And then loop it:

for i in Squares(4, 20):
    print i,

which prompts:

4 9 16

EDIT: edited to match the square definition, not the previous square power (sorry :P). Added +1 to the range to match the question example of the OP.

Answer 2

You can simplify the arithmetic by using

(n + 1)**2 == n**2 + (2*n + 1)

Here's how to do that using a generator function:

import math

def squares(lo, hi):
    root = int(math.ceil(lo ** 0.5))
    num = root ** 2
    delta = 2 * root + 1
    while num <= hi:
        yield num
        num += delta
        delta += 2

print list(squares(4, 16))
print list(squares(5, 50))
print list(squares(20, 90))

output

[4, 9, 16]
[9, 16, 25, 36, 49]
[25, 36, 49, 64, 81]

Here's an equivalent iterator class. I've given it a __repr__ method so it looks nice if you print an instance of this class.

import math

class Squares(object):
    def __init__(self, start, stop):
        self.start = start
        self.stop = stop
        root = int(math.ceil(start ** 0.5))
        self.num = root ** 2
        self.delta = 2 * root + 1

    def __repr__(self):
        return 'Squares(%d, %d)' % (self.start, self.stop)

    def __iter__(self): 
        return self

    def next(self):
        num = self.num
        if num > self.stop:
            raise StopIteration
        self.num += self.delta
        self.delta += 2
        return num

sq = Squares(4, 16)
print sq
for i in sq:
    print i

print list(Squares(5, 50))
print list(Squares(20, 90))

output

Squares(4, 16)
4
9
16
[9, 16, 25, 36, 49]
[25, 36, 49, 64, 81]

For Python 3, replace the next method name with __next__.

The usual Python range convention is to stop before you reach the high limit. To make this code comply with that convention, in the squares() generator change

while num <= hi:

to

while num < hi:

and in the Squares() class, change

if num > self.stop:

to

if num >= self.stop:

Answer 3

I think you meant to keep incrementing start until it reaches the next square, not to increment it only if it is a square:

def __next__(self):
    self.start += 1
    squareroot = math.sqrt(self.start)
    while squareroot != math.ceil(squareroot):
        if self.start > self.stop:
            raise StopIteration
        self.start += 1
        squareroot = math.sqrt(self.start)