Optimize custom_code

It would be awesome to optimize custom_code. Sometimes it works slowly.

Current implementation:

def custom_code(mask: str = '@###',
                char: str = '@', digit: str = '#') -> str:
    code = ''
    for p in mask:
        if p == char:
            code += choice(string.ascii_uppercase)
        elif p == digit:
            code += str(randint(0, 9))
        else:
            code += p
    return code

[lk-geimfari]

@duckyou Not yet! I’ll do it by myself. Thanks!

[duckyou]

Faster than original implementation:

def _randint(a, b):
    b = b - a
    return int(random.random() * b) + a


def new_custom_code(mask: str = '@###', char: str = '@',
                     digit: str = '#') -> str:

    code = bytearray(
        len(mask)
    )
    mask = mask.encode()
    char = ord(char)
    digit = ord(digit)

    for i, p in enumerate(mask):
        if p == char:
            a = _randint(65, 91)
        elif p == digit:
            a = _randint(48, 58)
        else:
            a = p
        code[i] = a
    return code.decode()

Compare code:

import time
import random
import string


# This function from Fluent Python by Luciano Ramalho.
def clock(func):
    def clocked(*args):
        t0 = time.time()
        result = func(*args)
        elapsed = time.time() - t0
        name = func.__name__
        arg_str = ', '.join(repr(arg) for arg in args)
        print('[%0.8fs] %s(%s) -> %r' % (elapsed, name, arg_str, result))
        return result

    return clocked


def orig_custom_code(mask: str = '@###', char: str = '@',
                 digit: str = '#') -> str:
    code = ''
    for p in mask:
        if p == char:
            code += random.choice(string.ascii_uppercase)
        elif p == digit:
            code += str(random.randint(0, 9))
        else:
            code += p

    return code


def _randint(a, b):
    b = b - a
    return int(random.random() * b) + a


def new_custom_code(mask: str = '@###', char: str = '@',
                     digit: str = '#') -> str:

    code = bytearray(
        len(mask)
    )
    mask = mask.encode()
    char = ord(char)
    digit = ord(digit)

    for i, p in enumerate(mask):
        if p == char:
            a = _randint(65, 91)
        elif p == digit:
            a = _randint(48, 58)
        else:
            a = p
        code[i] = a
    return code.decode()


@clock
def test_orig(count):
    db = [orig_custom_code('@_#_@@##+@_#_@@##+') for _ in range(count)]
    return 'Generated {}'.format(count)


@clock
def test_new(count):
    db = [new_custom_code('@_#_@@##+@_#_@@##+') for _ in range(count)]
    return 'Generated {}'.format(count)


if __name__ == '__main__':
    test_orig(100000)
    test_new(100000)

Results:

[3.20863533s] test_orig(100000) -> 'Generated 100000'
[1.42319679s] test_new(100000) -> 'Generated 100000'