Python – Calculating the square numbers within a range (python)


I want to be able to execute the following code:

for i in Squares(5, 50):

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):

I would expect the output to be:


Best Solution

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.