Search code examples
pythoniteratorabstractiondate-range

Generalized range / count functions for any type with __add__

Basically, I would like to be able to simply do something like:

 from datetime import datetime, timedelta
 from itertools import count
 start, end = datetime(2017,1,1), datetime(2018,1,1)
 calendar = count(start, timedelta(days=1))
 # Or better
 calendar = range(start, end, timedelta(days=1))

Is there a simple way to achieve this using builting/stdlib functionality? It says it needs a number on input - I think it should easily abstract to anything you can use the + operator on, but it does not.

Solution

  • You could of course redefine count to accept any types that can be added. Otherwise if you don't want to define your own types you can use itertools to do the same:

     calendar = itertools.accumulate( 
              itertools.chain( [start],
                               itertools.repeat( timedelta(days=1) ) ) )