vicky-remi-wedding-website/france/france_rsvp/codes.py

import logging
import random
import sys

from .database import get_db


ALPHABET = '0123456789bcdefghjkmnpqrstuvwxyz'
assert len(ALPHABET) == 32

LENGTH = 4


CORRECT = {
    'i': '1',
    'l': '1',
    'o': '0',
}


def correct_code(code):
    code = code.lower()
    fixed_code = ''.join(CORRECT.get(c, c) for c in code)
    return fixed_code


def list_errors(code):
    for place in range(LENGTH):
        for replacement in ALPHABET:
            if replacement == code[place]:
                continue
            new_code = code[:place] + replacement + code[place + 1:]
            yield new_code


def main():
    logging.basicConfig(level=logging.INFO)

    if len(sys.argv) <= 1:
        number = 1
    elif len(sys.argv) == 2:
        number = int(sys.argv[1], 10)
    else:
        raise ValueError

    # Get the current codes
    with get_db() as db:
        correct_codes = set(row[0] for row in db.execute(
            '''\
            SELECT code FROM families;
            ''',
        ))

    # Augment them with possible errored codes
    errored_codes = set()
    for code in correct_codes:
        errored_codes.update(list_errors(code))
    assert len(errored_codes) == LENGTH * (len(ALPHABET) - 1) * len(correct_codes)

    # Generate new codes
    generated = 0
    while generated < number:
        code = ''.join(random.choice(ALPHABET) for _ in range(LENGTH))
        if code in correct_codes or code in errored_codes:
            continue
        correct_codes.add(code)
        errored_codes.update(list_errors(code))
        print(code)
        generated += 1