Source

intspan / intspan.py

Full commit

import sys
from itertools import groupby, count


_PY3 = sys.version_info[0] > 2
if _PY3:
    basestring = str
        
class intspan(set):
    def __init__(self, initial=None):
        set.__init__(self)
        if initial:
            self.update(initial)
        
    def update(self, items):
        set.update(self, self._parse_range(items))
        
    def discard(self, items):
        for item in self._parse_range(items):
            set.discard(self, item)
            
    def remove(self, items):
        for item in self._parse_range(items):
            set.remove(self, item)
            
    def add(self, items):
        for item in self._parse_range(items):
            set.add(self, item)
                
    def __iter__(self):
        """
        Iterate in ascending order.
        """
        return iter(sorted(set.__iter__(self)))
    
    @staticmethod
    def _parse_range(datum):
        if isinstance(datum, basestring):
            result = []
            for part in datum.split(','):
                if '-' in part:
                    start, stop = part.split('-')
                    result.extend(list(range(int(start), int(stop)+1)))
                else:
                    result.append(int(part))
            return result
        else:
            return datum if hasattr(datum, '__iter__') else [ datum ]
        
    @staticmethod
    def _as_range(iterable): 
        l = list(iterable)
        if len(l) > 1:
            return '{0}-{1}'.format(l[0], l[-1])
        else:
            return '{0}'.format(l[0])

    def __repr__(self):        
        return 'intspan({0!r})'.format(self.__str__())
    
    def __str__(self):
        items = sorted(self)
        return ','.join(self._as_range(g) for _, g in groupby(items, key=lambda n, c=count(): n-next(c)))

    # see Jeff Mercado's answer to http://codereview.stackexchange.com/questions/5196/grouping-consecutive-numbers-into-ranges-in-python-3-2
    # see also: http://stackoverflow.com/questions/2927213/python-finding-n-consecutive-numbers-in-a-list


# It might be interesting to have a metaclass factory that could create
# spansets of things other than integers. For example, enumerateds defined
# by giving a universe of possible options. Or characters.