Morse code trainer

morse.py

from pyaudio import PyAudio
import math
import random

frequency = 700
sample_rate = 22050
volume = 0.5
wpm = 30

element_length = int((sample_rate * 60) / (50 * wpm))

CODE = {'A': '.-',     'B': '-...',   'C': '-.-.',
        'D': '-..',    'E': '.',      'F': '..-.',
        'G': '--.',    'H': '....',   'I': '..',
        'J': '.---',   'K': '-.-',    'L': '.-..',
        'M': '--',     'N': '-.',     'O': '---',
        'P': '.--.',   'Q': '--.-',   'R': '.-.',
     	'S': '...',    'T': '-',      'U': '..-',
        'V': '...-',   'W': '.--',    'X': '-..-',
        'Y': '-.--',   'Z': '--..',

#        '0': '-----',  '1': '.----',  '2': '..---',
#        '3': '...--',  '4': '....-',  '5': '.....',
#        '6': '-....',  '7': '--...',  '8': '---..',
#        '9': '----.'
        }

def envelope(length):
    """Ramp up/down to prevent audio clicks."""
    # TODO: set a reasonable number of samples based on sample rate & wpm
    samples = 10
    ret = [t / samples for t in range(samples)]
    ret += [1] * (length - samples - samples)
    ret += [1 - (t/samples) for t in range(samples)]
    return ret

def tone(elements):
    """Generate a fixed tone for a given number of element lengths."""
    global element_length, frequency, sample_rate, volume
    s = lambda t: volume * math.sin(2 * math.pi * frequency * t / sample_rate)
    env = envelope(element_length * elements)
    return [int(s(t) * env[t] * 0x7f + 0x80) for t in range(element_length * elements)]

def dit():
    return tone(1)

def dah():
    return tone(3)

def space():
    global element_length
    return [0x80] * element_length

def morse(chars):
    ret = []
    for c in chars:
        if c == '.':
            ret += dit()
        elif c == '-':
            ret += dah()
        elif c == ' ':
            ret += space()*2
        ret += space()
    return ret

def play(p, samples, sample_rate):
    restframes = len(samples) % sample_rate
    stream = p.open(format=p.get_format_from_width(1),
                    channels=1,
                    rate=sample_rate,
                    output=True)
    stream.write(bytes(bytearray(samples)))

    stream.write(b'\x80' * restframes)

    stream.stop_stream()
    stream.close()

p = PyAudio()
while True:
    char = random.choice(list(CODE.keys()))
    play(p, morse(CODE[char]), sample_rate)
    user_input = input(':')
    if char == user_input.upper():
        print("Correct")
    else:
        print("Incorrect, it was '" + char + "'")

p.terminate()