powerpod-command

#!/usr/bin/python
import argparse
import datetime
import functools
import logging
import os
import os.path
import re
import simplejson
import sys
import time
import traceback
import urllib2

import powerpod

LOGGER = logging.getLogger(__name__)

CMD_SPLIT = re.compile(r'(?:[^\\\s]|\\.)+')
WHITESPACE_QUOTE = functools.partial(re.compile(r'(\\s)').sub, r'\\\1')

ACTIONS = {}

def add_action(cls):
	ACTIONS[cls.PARSER.prog] = cls

class Action(object):
	PARSER = NotImplemented
	def __init__(self, extra):
		self.extra = extra

@add_action
class GetAllRidesAction(Action):
	PARSER = argparse.ArgumentParser('get_all_rides', description='Fetch all rides into ride_directory')
	PARSER.add_argument('--no-clobber', dest='existing', action='store_const', const='no_clobber')
	PARSER.add_argument('--force', dest='existing', action='store_const', const='force')
	PARSER.add_argument(
			'--directory',
			dest='directory',
			default='./rides',
	)
	def run(self, protocol, args):
		response = protocol.do_command(powerpod.GetFileListCommand())
		LOGGER.debug(repr(response))
		extra = []
		if self.extra.existing == 'force':
			extra.append('--force')
		elif self.extra.existing == 'no_clobber':
			extra.append('--no-clobber')
		for i, header in enumerate(response.records):
			filepath = os.path.join(self.extra.directory, header.to_filename())
			args.actions.append(make_action("get_ride {} {} {}".format(i, WHITESPACE_QUOTE(filepath), ' '.join(extra))))

@add_action
class ListRidesAction(Action):
	PARSER = argparse.ArgumentParser('list_rides', description='Print information about all rides on the device to stdout')
	def run(self, protocol, args):
		response = protocol.do_command(powerpod.GetFileListCommand())
		LOGGER.debug(repr(response))
		for i, header in enumerate(response.records):
			sys.stdout.write("{} {}\n".format(i, header.to_filename()))

@add_action
class GetRideAction(Action):
	PARSER = argparse.ArgumentParser('get_ride')
	PARSER.add_argument('index', type=int)
	PARSER.add_argument('filename')
	PARSER.add_argument('--no-clobber', dest='existing', action='store_const', const='no_clobber')
	PARSER.add_argument('--force', dest='existing', action='store_const', const='force')
	def run(self, protocol, args):
		time.sleep(1)
		response = protocol.do_command(powerpod.GetFileCommand(self.extra.index))
		filename = self.extra.filename
		if filename is sys.stdout:
			out = sys.stdout
		else:
			if os.path.exists(filename):
				if self.extra.existing == 'force':
					pass
				else:
					if self.extra.existing is None:
						LOGGER.warning('Will not overwrite {!r}; use --force or --no-clobber'.format(filename))
					return
			out = open(filename, 'w')
		LOGGER.info("index=%s header=%s filename=%s", self.extra.index, response.ride_data.get_header(), filename)
		out.write(response.ride_data.to_binary())

@add_action
class EraseAllCommand(Action):
	PARSER = argparse.ArgumentParser('erase_all')
	def run(self, protocol, args):
		protocol.do_command(powerpod.EraseAllCommand())
		rides = protocol.do_command(powerpod.GetFileListCommand())
		assert not rides.size, rides

@add_action
class GetOdometerCommand(Action):
	PARSER = argparse.ArgumentParser('get_odometer')
	def run(self, protocol, args):
		response = protocol.do_command(powerpod.GetOdometerCommand())
		sys.stdout.write('{} km\n'.format(response.distance_km))

@add_action
class SetOdometerCommand(Action):
	PARSER = argparse.ArgumentParser('set_odometer')
	PARSER.add_argument('distance', type=float, help='km rounded to 1 decimal place')
	def run(self, protocol, args):
		distance = round(self.extra.distance, 1)
		if distance != self.extra.distance:
			LOGGER.warning('rounding to one decimal place')
		protocol.do_command(powerpod.SetOdometerCommand(distance))
		response = protocol.do_command(powerpod.GetOdometerCommand())
		assert response.distance_km == distance, response.distance_km

@add_action
class GetUnitsCommand(Action):
	PARSER = argparse.ArgumentParser('get_units')
	def run(self, protocol, args):
		response = protocol.do_command(powerpod.GetOdometerCommand())
		sys.stdout.write("{}\n".format(powerpod.SetUnitsCommand.LOOKUP[response.units_type]))

@add_action
class SetUnitsCommand(Action):
	PARSER = argparse.ArgumentParser('set_units')
	PARSER.add_argument('units_type', type=powerpod.SetUnitsCommand.LOOKUP.index)
	def run(self, protocol, args):
		protocol.do_command(powerpod.SetUnitsCommand(self.extra.units_type))
		response = protocol.do_command(powerpod.GetOdometerCommand())
		assert response.units_type == self.extra.units_type, response.units_type

@add_action
class GetFirmwareVersionCommand(Action):
	PARSER = argparse.ArgumentParser('get_firmware_version')
	def run(self, protocol, args):
		response = protocol.do_command(powerpod.GetFirmwareVersionCommand())
		sys.stdout.write('{} {}\n'.format(response.version_encoded, response.version))

@add_action
class GetFirmwareVersionCommand(Action):
	PARSER = argparse.ArgumentParser('check_firmware_version', description="Exits with status 1 if not latest")
	PARSER.add_argument('type', choices=('PwrPod', 'iBike'))
	def run(self, protocol, args):
		response = protocol.do_command(powerpod.GetFirmwareVersionCommand())
		# SUPER PORCELAIN!
		# I think I'm actually meant to read FWstatus.txt, but that is huge and
		# clearly includes much more than the latest version!
		resp = urllib2.urlopen('http://ibikesports.com/iBike_update/')
		resp_content = resp.read()
		versions = re.compile(r'<a href="(PwrPod|iBike)(\d+)\.iBFW3.txt">').findall(resp_content)
		best_version = max(int(version) for type, version in versions if type == self.extra.type)
		if best_version != response.version_encoded:
			sys.stdout.write('{!r} is not latest; {!r} is available\n'.format(response.version_encoded, best_version))
			sys.exit(1)

@add_action
class SetTrainerWeightsCommand(Action):
	PARSER = argparse.ArgumentParser('set_trainer_weights', description="""All coefficients are in terms of a polynomial in mph, outputting Watts.""")
	PARSER.add_argument('constant', type=float)
	PARSER.add_argument('linear', type=float)
	PARSER.add_argument('quadratic', type=float)
	PARSER.add_argument('cubic', type=float)
	def run(self, protocol, args):
		command = powerpod.SetTrainerWeightsCommand(
				self.extra.constant,
				self.extra.linear,
				self.extra.quadratic,
				self.extra.cubic,
		)
		response = protocol.do_command(command)

def interval_format(string):
	power_watts, work_secs, rest_secs = map(int, string.split(':'))
	return powerpod.NewtonInterval(power_watts, work_secs, rest_secs)

@add_action
class SetIntervalsCommand(Action):
	PARSER = argparse.ArgumentParser('set_intervals')
	PARSER.add_argument('intervals', type=interval_format, nargs='*')
	def run(self, protocol, args):
		protocol.do_command(powerpod.SetIntervalsCommand(
						len(self.extra.intervals),
						len(self.extra.intervals),
						self.extra.intervals,
		))

@add_action
class GetSerialNumberCommand(Action):
	PARSER = argparse.ArgumentParser('get_serial_number')
	def run(self, protocol, args):
		response = protocol.do_command(powerpod.GetSerialNumberCommand())
		sys.stdout.write('{}\n'.format(response.as_hex))

@add_action
class GetSpaceUsageCommand(Action):
	PARSER = argparse.ArgumentParser('get_space_used')
	def run(self, protocol, args):
		response = protocol.do_command(powerpod.GetSpaceUsageCommand())
		sys.stdout.write('{} %\n'.format(response.used_percentage))

TIME_FORMAT = lambda string: datetime.datetime.strptime(string, '%Y-%m-%dT%H:%M:%S')

@add_action
class SetTimeCommand(Action):
	PARSER = argparse.ArgumentParser('set_time')
	PARSER.add_argument('--time', type=TIME_FORMAT, help='eg. 2016-01-31T22:01:12; default=local time', required=False)
	def run(self, protocol, args):
		if self.extra.time is None:
			time = datetime.datetime.now()
		else:
			time = self.extra.time
		protocol.do_command(powerpod.SetTimeCommand(0, powerpod.NewtonTime.from_datetime(time)))

@add_action
class GetDefaultProfileCommand(Action):
	PARSER = argparse.ArgumentParser('get_default_profile')
	PARSER.add_argument('number', choices=(0, 1, 2, 3), type=int)
	def run(self, protocol, args):
		response = protocol.do_command(powerpod.GetProfileNumberCommand())
		sys.stdout.write("{}\n".format(response.number))

@add_action
class SetDefaultProfileCommand(Action):
	PARSER = argparse.ArgumentParser('set_default_profile')
	PARSER.add_argument('number', choices=(0, 1, 2, 3), type=int)
	def run(self, protocol, args):
		protocol.do_command(powerpod.SetProfileNumberCommand(self.extra.number))
		response = protocol.do_command(powerpod.GetProfileNumberCommand())
		assert self.extra.number == response.number, response.number

@add_action
class DumpProfilesCommand(Action):
	""" Since the profile editing commands are a different shape to profile getting, there seems little point in a raw dump. """
	PARSER = argparse.ArgumentParser('dump_profiles')
	PARSER.add_argument('--number', choices=(0, 1, 2, 3), type=int)
	def run(self, protocol, args):
		response = protocol.do_command(powerpod.GetProfileDataCommand())
		if self.extra.number is not None:
			data = simplejson.dumps(response.records[self.extra.number]._asdict())
		else:
			data = simplejson.dumps([profile._asdict() for profile in response.records])
		sys.stdout.write("{}\n".format(data))

@add_action
class RestoreProfilesCommand(Action):
	PARSER = argparse.ArgumentParser('restore_profiles', description="Note that this temporarily sets the default profile, and not all data is restored. Input should be JSON from dump_profiles.")
	PARSER.add_argument('--number', choices=(0, 1, 2, 3), type=int)
	def run(self, protocol, args):
		response = protocol.do_command(powerpod.GetProfileNumberCommand())
		original_profile = current_profile = response.number
		data = simplejson.load(sys.stdin)
		if self.extra.number is not None:
			profiles = [(self.extra.number, data)]
		else:
			profiles = list(enumerate(data))
		try:
			for number, profile in profiles:
				if current_profile != number:
					current_profile = number
					protocol.do_command(powerpod.SetProfileNumberCommand(number))
				params = []
				for field in powerpod.SetProfileDataCommand._fields:
					params.append(profile.pop(field))
				protocol.do_command(powerpod.SetProfileDataCommand(*params))
				params = []
				for field in powerpod.SetProfileData2Command._fields:
					params.append(profile.pop(field))
				protocol.do_command(powerpod.SetProfileData2Command(*params))
				if profile:
					LOGGER.warn("left over profile %s data: %r", number, profile)
		except Exception:
			traceback.print_exc()
			if current_profile != original_profile:
				protocol.do_command(powerpod.SetProfileNumberCommand(original_profile))
			sys.exit(1)
		if current_profile != original_profile:
			protocol.do_command(powerpod.SetProfileNumberCommand(original_profile))

@add_action
class UpdateProfileCommand(Action):
	PARSER = argparse.ArgumentParser('update_profile', description="Note that this temporarily sets the default profile.")
	PARSER.add_argument('number', choices=(0, 1, 2, 3), type=int)
	for command in (powerpod.SetProfileDataCommand, powerpod.SetProfileData2Command):
		for field, typ in zip(command._fields, command.SHAPE[1:]):
			if typ in ('h', 'H'):
				field_type = int
			else:
				assert typ == 'f', typ
				field_type = float
			PARSER.add_argument('--' + field.replace('_', '-'), type=field_type)
	def run(self, protocol, args):
		response = protocol.do_command(powerpod.GetProfileNumberCommand())
		original_profile = response.number
		if original_profile != self.extra.number:
			protocol.do_command(powerpod.SetProfileNumberCommand(self.extra.number))
		profile = protocol.do_command(powerpod.GetProfileDataCommand()).records[self.extra.number]
		coalesce = lambda field: getattr(self.extra, field) if getattr(self.extra, field) is not None else getattr(profile, field)
		try:
			params = []
			used = False
			for field in powerpod.SetProfileDataCommand._fields:
				if getattr(self.extra, field) is None:
					params.append(getattr(profile, field))
				else:
					used = True
					params.append(getattr(self.extra, field))
			if used:
				protocol.do_command(powerpod.SetProfileDataCommand(*params))
			params = []
			used = False
			for field in powerpod.SetProfileData2Command._fields:
				if getattr(self.extra, field) is None:
					params.append(getattr(profile, field))
				else:
					used = True
					params.append(getattr(self.extra, field))
			if used:
				protocol.do_command(powerpod.SetProfileData2Command(*params))
		except Exception:
			traceback.print_exc()
			if self.extra.number != original_profile:
				protocol.do_command(powerpod.SetProfileNumberCommand(original_profile))
			sys.exit(1)
		if self.extra.number != original_profile:
			protocol.do_command(powerpod.SetProfileNumberCommand(original_profile))

@add_action
class DumpScreensCommand(Action):
	""" Since the profile editing commands are a different shape to profile getting, there seems little point in a raw dump. """
	PARSER = argparse.ArgumentParser('dump_screens')
	PARSER.add_argument('--number', choices=(0, 1, 2, 3), type=int)
	def run(self, protocol, args):
		response = protocol.do_command(powerpod.GetAllScreensCommand())
		if self.extra.number is not None:
			data = simplejson.dumps(response.records[self.extra.number].to_dict())
		else:
			data = simplejson.dumps([profile.to_dict() for profile in response.records])
		sys.stdout.write("{}\n".format(data))

@add_action
class RestoreScreensCommand(Action):
	PARSER = argparse.ArgumentParser('restore_screens', description="Note that this temporarily sets the default profile, and not all data is restored. Input should be JSON from dump_screens.")
	PARSER.add_argument('--number', choices=(0, 1, 2, 3), type=int)
	def run(self, protocol, args):
		response = protocol.do_command(powerpod.GetProfileNumberCommand())
		original_profile = current_profile = response.number
		data = simplejson.load(sys.stdin)
		if self.extra.number is not None:
			screenss = [(self.extra.number, data)]
		else:
			screenss = list(enumerate(data))
		try:
			for number, screens in screenss:
				if current_profile != number:
					current_profile = number
					protocol.do_command(powerpod.SetProfileNumberCommand(number))
				protocol.do_command(powerpod.SetScreensCommand(powerpod.NewtonProfileScreens.from_dict(screens)))
		except Exception:
			traceback.print_exc()
			if current_profile != original_profile:
				protocol.do_command(powerpod.SetProfileNumberCommand(original_profile))
			sys.exit(1)
		if current_profile != original_profile:
			protocol.do_command(powerpod.SetProfileNumberCommand(original_profile))

def make_action(string):
	parts = CMD_SPLIT.findall(string)
	if not parts:
		raise ValueError("Invalid command {}".format(string))
	if parts[0] not in ACTIONS:
		raise ValueError("Unknown command {}".format(string))
	action_class = ACTIONS[parts[0]]
	args = action_class.PARSER.parse_args(parts[1:])
	return action_class(args)

class HelpActions(argparse.Action):
	def __init__(self, *args, **kwargs):
		kwargs['nargs'] = 0
		super(HelpActions, self).__init__(*args, **kwargs)
	@staticmethod
	def __call__(parser, namespace, values, option_string):
		for action_class in ACTIONS.values():
			action_class.PARSER.print_help()
			sys.stdout.write('\n\n')
		sys.exit(0)

def arg_parser():
	parser = argparse.ArgumentParser()
	parser.add_argument('--port', default='/dev/ttyUSB0')
	parser.add_argument('--debug', default=False)
	parser.add_argument(
			'actions',
			nargs='+',
			type=make_action,
			help="""
				Actions to perform. See --help-actions for a list.
				Parameters to actions must be quoted with the action.
			""",
	)
	parser.add_argument(
			'--help-actions',
			action=HelpActions,
	)
	return parser

def main():
	args = arg_parser().parse_args()
	if args.debug:
		log_level = logging.DEBUG
	else:
		log_level = logging.INFO
	logging.basicConfig(level=log_level)
	kwargs = {}
	serial_connection = powerpod.NewtonSerialConnection(port=args.port)
	protocol = powerpod.NewtonSerialProtocol(serial_connection, device_side=False)
	for action in args.actions:
		action.run(protocol, args)

if __name__ == '__main__':
	main()