str_generator

str_generator.py

from itertools import product


def generate_from_pattern(pattern, open_sep='{{', close_sep='}}'):
    """ Generates all possible variations that fits pattern.

    For example:
    >>> generate_from_pattern('number {{9-11}}')
    ['number 9', 'number 10', 'number 11']
    >>> generate_from_pattern('number {{1-3,11}}')
    ['number 1', 'number 2', 'number 3', 'number 11']
    >>> generate_from_pattern('letter {{a-c,e}}')
    ['letter a', 'letter b', 'letter c', 'letter e']
    >>> generate_from_pattern('{{a-c}}{{1-3}}')
    ['a1', 'a2', 'a3', 'b1', 'b2', 'b3', 'c1', 'c2', 'c3']
    >>> generate_from_pattern('just do nothing')
    ['just do nothing']
    """

    tokens = _tokenize(pattern, open_sep, close_sep)
    result = _generate(tokens, open_sep, close_sep)
    return result


def _tokenize(s, open_sep='{{', close_sep='}}'):
    """
    >>> _tokenize('aa {{a,b,c}} z')
    ['aa ', '{{', 'a,b,c', '}}', ' z']
    """

    tokens = []
    for x in s.split(open_sep):
        for y in x.split(close_sep):
            tokens.append(y)
            tokens.append(close_sep)
        tokens.pop()
        tokens.append(open_sep)
    tokens.pop()
    return tokens


def _generate_one(s):
    """
    >>> _generate_one("a-c,1-3,4,zaz,123")
    ['a', 'b', 'c', '1', '2', '3', '4', 'zaz', '123']
    """
    groups = s.split(',')
    res = []
    for group in groups:
        if '-' in group:
            frm, to = group.split('-')
        else:
            frm = to = group

        try:
            frm, to = int(frm), int(to)
            res.extend(str(c) for c in range(frm, to + 1))
        except ValueError:
            if len(frm) == 1:
                res.extend(chr(c) for c in range(ord(frm), ord(to) + 1))
            else:
                res.append(frm)

    return res


def _generate(tokens, open_sep='{{', close_sep='}}'):
    """
    >>> _generate(['x', '{{', '1-2', '}}', 'z'])
    ['x1z', 'x2z']
    """
    state = 'out'
    parsed = []
    for token in tokens:
        if token == open_sep:
            state = 'in'
        elif token == close_sep:
            state = 'out'
        else:
            if state == 'in':
                parsed.append(_generate_one(token))
            else:
                parsed.append([token])

    return list(''.join(item) for item in product(*parsed))