wg1220/wg1220.py

import serial
import operator
import logging
import struct
import math
import argparse


log = logging.getLogger(__name__)


def checksum(msg):
    return reduce(operator.xor, (ord(x) for x in msg))


class Generator(object):

    def __init__(self):
        self.conn = serial.Serial(
                port='/dev/ttyUSB0',
                baudrate=57600,
                bytesize=serial.EIGHTBITS,
                parity=serial.PARITY_NONE,
                stopbits=serial.STOPBITS_ONE,
                timeout=1,
            )

    def send_command(self, cmd):
        if len(cmd) != 5:
            raise ValueError("cmd must be five characters")

        def fmt(msg):
            return ' '.join(hex(ord(x)) for x in msg)

        message = 'X{}{}'.format(cmd, chr(checksum(cmd)))
        log.debug("send {}".format(fmt(message)))

        self.conn.write(message)
        result = self.conn.read(7)

        if len(result) == 0:
            raise IOError("No result returned")

        log.debug("recv {}".format(fmt(result)))

        if checksum(result[1:6]) != ord(result[6]):
            raise IOError("Wrong checksum")

        return result[2:5]

    def set_remote(self, enable):
        self.send_command('R{}xxx'.format(chr(int(not enable))))
        r = self.send_command('rxxxx')

    @property
    def remote_enabled(self):
        r = self.send_command('rxxxx')
        return r[0] == '0'

    @property
    def frequency(self):
        r = self.send_command('fxxxx')
        hi, lo, exp = struct.unpack('BBB', r)
        return ((hi << 8) | lo) / 1000.0 * 10**exp

    @frequency.setter
    def frequency(self, freq):
        exp = int(math.log(freq, 10))
        mantissa = int(round(freq / math.exp(exp * math.log(10)))) * 1000
        hi, lo = mantissa >> 8, mantissa & 0xff 
        self.send_command('F{}{}{}{}'.format(chr(hi), chr(lo), chr(exp), chr(0xff)))

    @property
    def amplitude(self):
        r = self.send_command('axxxx')
        hi, lo = struct.unpack('BB', r[:2])
        return ((hi << 8) | lo) / 1000.0

    @amplitude.setter
    def amplitude(self, amp):
        val = int(amp * 1000)
        hi, lo = val >> 8, val & 0xff
        self.send_command('A{}{}x{}'.format(chr(hi), chr(lo), chr(0xff)))

    @property
    def offset(self):
        r = self.send_command('oxxxx')
        hi, lo = struct.unpack('BB', r[:2])
        return ((hi << 8) | lo) / 1000.0

    @offset.setter
    def offset(self, off):
        val = int(off * 1000)
        hi, lo = val >> 8, val & 0xff
        self.send_command('O{}{}x{}'.format(chr(hi), chr(lo), chr(0xff)))

    @property
    def dutycycle(self):
        r = self.send_command('dxxxx')
        return ord(r[0])

    @dutycycle.setter
    def dutycycle(self, cycle):
        if cycle < 0 or cycle > 100:
            raise ValueError("cycle must be between 0 and 100")

        self.send_command('D{}xx{}'.format(chr(cycle), chr(0xff)))


def main():
    logging.basicConfig(level=logging.INFO, format="[%(asctime)s]  %(message)s")
    g = Generator()

    parser = argparse.ArgumentParser()
    parser.add_argument('-f', '--frequency', type=float, default=None)
    parser.add_argument('-a', '--amplitude', type=float, default=None)
    parser.add_argument('-o', '--offset', type=float, default=None)
    parser.add_argument('-d', '--dutycycle', type=float, default=None)
    args = parser.parse_args()

    g.set_remote(True)

    if args.frequency:
        g.frequency = args.frequency

    if args.amplitude:
        g.amplitude = args.amplitude

    if args.offset:
        g.offset = args.offset

    if args.dutycycle:
        g.dutycycle = argsycle

    print("Frequency: {} Hz".format(g.frequency))
    print("Amplitude: {} V".format(g.amplitude))
    print("Offset: {} V".format(g.offset))
    print("Duty cycle: {} %".format(g.dutycycle))

    g.set_remote(False)
    

if __name__ == '__main__':
    main()