ODrive/tools/odrivetool

#!/usr/bin/env python3
"""
ODrive command line utility
"""

from __future__ import print_function
import sys
import os
import argparse
import time

sys.path.insert(0, os.path.join(os.path.dirname(os.path.dirname(
                    os.path.realpath(__file__))),
                    "Firmware", "fibre", "python"))
import fibre.discovery
from fibre import Logger, Event
import odrive
from odrive.utils import OperationAbortedException
from odrive.configuration import *

# Flush stdout by default
# Source:
# https://stackoverflow.com/questions/230751/how-to-flush-output-of-python-print
old_print = print
def print(*args, **kwargs):
    kwargs.pop('flush', False)
    old_print(*args, **kwargs)
    file = kwargs.get('file', sys.stdout)
    file.flush() if file is not None else sys.stdout.flush()

script_path=os.path.dirname(os.path.realpath(__file__))

## Parse arguments ##
parser = argparse.ArgumentParser(description='ODrive command line utility\n'
                                             'Running this tool without any arguments is equivalent to running `odrivetool shell`\n',
                                 formatter_class=argparse.RawTextHelpFormatter)

# Subcommands
subparsers = parser.add_subparsers(help='sub-command help', dest='command')
shell_parser = subparsers.add_parser('shell', help='Drop into an interactive python shell that lets you interact with the ODrive(s)')
shell_parser.add_argument("--no-ipython", action="store_true",
                          help="Use the regular Python shell "
                          "instead of the IPython shell, "
                          "even if IPython is installed.")

dfu_parser = subparsers.add_parser('dfu', help="Upgrade the ODrive device firmware."
                                               "If no serial number is specified, the first ODrive that is found is updated")
dfu_parser.add_argument('file', metavar='HEX', nargs='?',
                        help='The .hex file to be flashed. Make sure target board version '
                        'of the firmware file matches the actual board version. '
                        'You can download the latest release manually from '
                        'https://github.com/madcowswe/ODrive/releases. '
                        'If no file is provided, the script automatically downloads '
                        'the latest firmware.')


dfu_parser = subparsers.add_parser('backup-config', help="Saves the configuration of the ODrive to a JSON file")
dfu_parser.add_argument('file', nargs='?',
                        help="Path to the file where to store the data. "
                        "If no path is provided, the configuration is stored in {}.".format(tempfile.gettempdir()))

dfu_parser = subparsers.add_parser('restore-config', help="Restores the configuration of the ODrive from a JSON file")
dfu_parser.add_argument('file', nargs='?',
                        help="Path to the file that contains the configuration data. "
                        "If no path is provided, the configuration is loaded from {}.".format(tempfile.gettempdir()))

code_generator_parser = subparsers.add_parser('generate-code', help="Process a jinja2 template, passing the ODrive's JSON data as data input")
code_generator_parser.add_argument("-t", "--template", type=argparse.FileType('r'),
                    help="the code template")
code_generator_parser.add_argument("-o", "--output", type=argparse.FileType('w'), default='-',
                    help="path of the generated output")
code_generator_parser.set_defaults(template = os.path.join(script_path, 'odrive_header_template.h.in'))

subparsers.add_parser('liveplotter', help="For plotting of odrive parameters (i.e. position) in real time")
subparsers.add_parser('drv-status', help="Show status of the on-board DRV8301 chips (for debugging only)")
subparsers.add_parser('rate-test', help="Estimate the average transmission bandwidth over USB")
subparsers.add_parser('udev-setup', help="Linux only: Gives users on your system permission to access the ODrive by installing udev rules")

# General arguments
parser.add_argument("-p", "--path", metavar="PATH", action="store",
                    help="The path(s) where ODrive(s) should be discovered.\n"
                    "By default the script will connect to any ODrive on USB.\n\n"
                    "To select a specific USB device:\n"
                    "  --path usb:BUS:DEVICE\n"
                    "usbwhere BUS and DEVICE are the bus and device numbers as shown in `lsusb`.\n\n"
                    "To select a specific serial port:\n"
                    "  --path serial:PATH\n"
                    "where PATH is the path of the serial port. For example \"/dev/ttyUSB0\".\n"
                    "You can use `ls /dev/tty*` to find the correct port.\n\n"
                    "You can combine USB and serial specs by separating them with a comma (no space!)\n"
                    "Example:\n"
                    "  --path usb,serial:/dev/ttyUSB0\n"
                    "means \"discover any USB device or a serial device on /dev/ttyUSB0\"")
parser.add_argument("-s", "--serial-number", action="store",
                    help="The 12-digit serial number of the device. "
                         "This is a string consisting of 12 upper case hexadecimal "
                         "digits as displayed in lsusb. \n"
                         "    example: 385F324D3037\n"
                         "You can list all devices connected to USB by running\n"
                         "(lsusb -d 1209:0d32 -v; lsusb -d 0483:df11 -v) | grep iSerial\n"
                         "If omitted, any device is accepted.")
parser.add_argument("-v", "--verbose", action="store_true",
                    help="print debug information")
parser.add_argument("--version", action="store_true",
                    help="print version information and exit")

parser.set_defaults(path="usb")
args = parser.parse_args()

# Default command
if args.command is None:
    args.command = 'shell'
    args.no_ipython = False

logger = Logger(verbose=args.verbose)

def print_version():
    sys.stderr.write("ODrive control utility v" + odrive.__version__ + "\n")
    sys.stderr.flush()

app_shutdown_token = Event()

try:
    if args.version == True:
        print_version()

    elif args.command == 'shell':
        print_version()
        if ".dev" in odrive.__version__:
            print("")
            logger.warn("Developer Preview")
            print("  If you find issues, please report them")
            print("  on https://github.com/madcowswe/ODrive/issues")
            print("  or better yet, submit a pull request to fix it.")
            print("")
        import odrive.shell
        odrive.shell.launch_shell(args, logger, app_shutdown_token)

    elif args.command == 'dfu':
        print_version()
        import odrive.dfu
        odrive.dfu.launch_dfu(args, logger, app_shutdown_token)

    elif args.command == 'liveplotter':
        from odrive.utils import start_liveplotter
        print("Waiting for ODrive...")
        my_odrive = odrive.find_any(path=args.path, serial_number=args.serial_number,
                                              search_cancellation_token=app_shutdown_token,
                                              channel_termination_token=app_shutdown_token)

        # If you want to plot different values, change them here.
        # You can plot any number of values concurrently.
        cancellation_token = start_liveplotter(lambda: [
            my_odrive.axis0.encoder.pos_estimate,
            my_odrive.axis1.encoder.pos_estimate,
        ])

        print("Showing plot. Press Ctrl+C to exit.")
        while not cancellation_token.is_set():
            time.sleep(1)

    elif args.command == 'drv-status':
        from odrive.utils import print_drv_regs
        print("Waiting for ODrive...")
        my_odrive = odrive.find_any(path=args.path, serial_number=args.serial_number,
                                              search_cancellation_token=app_shutdown_token,
                                              channel_termination_token=app_shutdown_token)
        print_drv_regs("Motor 0", my_odrive.axis0.motor)
        print_drv_regs("Motor 1", my_odrive.axis1.motor)

    elif args.command == 'rate-test':
        from odrive.utils import rate_test
        print("Waiting for ODrive...")
        my_odrive = odrive.find_any(path=args.path, serial_number=args.serial_number,
                                              search_cancellation_token=app_shutdown_token,
                                              channel_termination_token=app_shutdown_token)
        rate_test(my_odrive)

    elif args.command == 'udev-setup':
        from odrive.utils import setup_udev_rules
        setup_udev_rules(logger)
    
    elif args.command == 'generate-code':
        from odrive.code_generator import generate_code
        my_odrive = odrive.find_any(path=args.path, serial_number=args.serial_number,
            channel_termination_token=app_shutdown_token)
        generate_code(my_odrive, args.template, args.output)

    elif args.command == 'backup-config':
        from odrive.configuration import backup_config
        print("Waiting for ODrive...")
        my_odrive = odrive.find_any(path=args.path, serial_number=args.serial_number,
                                              search_cancellation_token=app_shutdown_token,
                                              channel_termination_token=app_shutdown_token)
        backup_config(my_odrive, args.file, logger)

    elif args.command == 'restore-config':
        from odrive.configuration import restore_config
        print("Waiting for ODrive...")
        my_odrive = odrive.find_any(path=args.path, serial_number=args.serial_number,
                                              search_cancellation_token=app_shutdown_token,
                                              channel_termination_token=app_shutdown_token)
        restore_config(my_odrive, args.file, logger)

    else:
        raise Exception("unknown command: " + args.command)

except OperationAbortedException:
    logger.info("Operation aborted.")
finally:
    app_shutdown_token.set()