Iterator and Generator versions of Python's range()

Question

I have created iterator and generator versions of Python's range():

Generator version

def irange(*args):
 if len(args) > 3:
     raise TypeError('irange() expected at most 3 arguments, got %s' % (len(args)))
 elif len(args) == 1:
     start_element = 0
     end_element = args[0]
     step = 1
 else:
     start_element = args[0]
     end_element = args[1]
     if len(args) == 2:
         step = 1
     elif (args[2] % 1 == 0 and args[2] != 0):
         step = args[2]
     else:
         raise ValueError('irange() step argument must not be zero')

 if((type(start_element) is str) or (type(end_element) is str) 
    or (type(step) is str)):
     raise TypeError('irange() integer expected, got str')

 count = 0
 while (( start_element + step < end_element ) 
        if 0 < step else 
        ( end_element < start_element + step )) :
     if count == 0:
         item = start_element
     else:
         item = start_element + step
     start_element = item
     count +=1
     yield item

Iterator version

class Irange:
 def __init__(self, start_element, end_element=None, step=1):
     if step == 0:
         raise ValueError('Irange() step argument must not be zero')
     if((type(start_element) is str) or (type(end_element) is str) 
        or (type(step) is str)):
         raise TypeError('Irange() integer expected, got str')

     self.start_element = start_element
     self.end_element = end_element
     self.step = step
     self.index = 0

     if end_element is None:
        self.start_element = 0
        self.end_element = start_element

 def __iter__(self):
     return self

 def next(self):
     if self.index == 0:
         self.item = self.start_element
     else:
         self.item = self.start_element + self.step
     if self.step > 0:
         if self.item >= self.end_element:
                raise StopIteration
     elif self.step < 0:
         if self.item <= self.end_element:
             raise StopIteration

     self.start_element = self.item
     self.index += 1
     return self.item

Usage

 >>> for i in irange(2,5):
 ...  print i,
 2 3 4
 >>> for i in irange(2,-3,-1):
 ...  print i,
 2 1 0 -1 -2
 >>> for i in Irange(3):
 ...  print i,
 0 1 2

I would like to know if the approach is correct.

Answer 1

First of all, to validate that the function is working, it's good to use assert statements:

assert [0, 1, 2, 3, 4] == [x for x in irange(5)]
assert [2, 3, 4] == [x for x in irange(2, 5)]
assert [2, 1, 0, -1, -2] == [x for x in irange(2, -3, -1)]

With these statements covering my back, I refactored your irange method to this:

def irange(*args):
    len_args = len(args)

    if len_args > 3:
        raise TypeError('irange() expected at most 3 arguments, got %s' % len_args)

    if len_args < 1:
        raise TypeError('irange() expected at least 1 arguments, got %s' % len_args)

    sanitized_args = [int(x) for x in args]

    if len_args == 1:
        start_element = 0
        end_element = sanitized_args[0]
        step = 1
    else:
        start_element = sanitized_args[0]
        end_element = sanitized_args[1]
        step = 1 if len_args == 2 else sanitized_args[2]

    current = start_element

    if step > 0:
        def should_continue():
            return current < end_element
    else:
        def should_continue():
            return current > end_element

    while should_continue():
        yield current
        current += step

Points of improvement:

• Since len(args) is used repeatedly, I cache it in len_args
• Added len(args) < 1 check too, in the same fashion as len(args) > 3
• Simplified the type checking of args:
  • Sanitize with a single, simple list comprehension
  • If there are any non-integer arguments, a ValueError will be raised with a reasonably understandable error message
• Simplified the initialization of start_element, end_element and step
• Greatly simplified the stepping logic

As for the iterator version, it would be easiest and best to implement that in terms of the generator version.