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.

• Adam Wagner: how is this different than xrange? Or are you just testing your ability to write generators/iterators?
• Kinjal: This is same as xrange but using generator/iterator
• Adam Wagner: If you are looking for the behavior of an iterator, does this accomplish the same for you? def irange(*args): return iter(xrange(*args))
• Kinjal: this sounds good.
• P3trus: In python 2 xrange is an iterable returning an Iterator stackoverflow.com/a/10776268/639650

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.