dec8/dtk/argv.py

#! /usr/bin/env python3
# encoding: utf-8

import sys

from . import aux

class Action:
	"""
	A callable with signature:
	
	p_args is a list of (name, type, desc) tuples
	kv_args is a dict with {name: (type, default value, desc), ...}
	"""
	
	p_args = []
	kv_args = {}
	
	def __call__(self, cmd, **kwargs):
		"""
		This is passed the invoking command and all the args.
		
		Override this.
		"""
		
		raise NotImplementedError

class UnknownCommand(Exception):
	def __init__(self, cmd):
		super().__init__(" ".join(cmd))

class BadArg(Exception):
	def __init__(self, name, T, raw):
		super().__init__("unable to interpret %s as %s (%s)" %(raw, name, T.__name__))

class BadKeywordArg(Exception):
	def __init__(self, name):
		super().__init__("unknown keyword argument %s" %(name))

class IncompleteInvocation(Exception):
	def __init__(self):
		super().__init__("one or more mandatory arguments are missing")

class Invocation:
	"""
	A simple command line parser.
	
	Initialized with a (nested) routing dict which maps commands to Actions.
	
	COMMANDS ARGS [KV_ARGS]
	
	COMMANDS is one or more words describing the Action to take.
	
	ARGS is exactly N ordered words (N depends on the selected Action).
	
	KV_ARGS is zero or more name=value pairs. Their meaning depends on the selected Action.
	
	It is an error if any token remains after parsing.
	"""
	
	def __init__(self, routing):
		self.routing = routing
		
		self.action = None
		self.cmd = []
		self.args = {} # name: value
	
	def parse_argv(self, argv=None):
		"""
		Parses the supplied argv in order to pre-select an Action to call.
		
		If no argv is passed, uses sys.argv[1:].
		"""
		
		argv = argv or sys.argv[1:]
		cursor = self.routing
		consumed = 0
		
		# decode the command
		for arg in argv:
			self.cmd.append(arg)
			consumed += 1
			
			if arg in cursor:
				cursor = cursor[arg]
			else:
				raise UnknownCommand(self.cmd)
			
			if issubclass(type(cursor), Action):
				self.action = cursor
				break
		
		# at this point the expected layout of the remaining args is known
		
		if not self.action:
			return
		
		# pre-populate args
		for name, *_ in self.action.p_args:
			self.args[name] = None
		
		for name, (T, default, _) in self.action.kv_args.items():
			self.args[name] = T(default)
		
		# process positional (mandatory) args
		count_p_args = len(self.action.p_args)
		raw_p_args = argv[consumed:][:count_p_args]
		
		if not len(raw_p_args) == count_p_args:
			raise IncompleteInvocation
		
		for raw, (name, T, _) in zip(raw_p_args, self.action.p_args):
			try:
				self.args[name] = T(raw)
			except:
				raise BadArg(name, T, raw)
		
		consumed += count_p_args
		
		# process keyword (optional) args
		raw_kv_args = argv[consumed:]
		
		for arg in raw_kv_args:
			if "=" in arg:
				aname, avalue = arg.split("=", 1)
				
				if aname in self.action.kv_args:
					self.args[aname] = self.action.kv_args[aname][0](avalue)
				else:
					raise BadKeywordArg(arg)
			else:
				raise BadKeywordArg(arg)
	
	def execute(self):
		self.action(self.cmd, **self.args)
	
	def help(self, stream=sys.stdout):
		"""
		Prints auto-generated usage information.
		"""
		
		print("Usage:")
		
		for cmd, act in aux.flatten_dict(self.routing):
			synopsis = ["$", sys.argv[0], cmd]
			args = []
			
			for name, T, desc in act.p_args:
				synopsis.append(name)
				args.append("* %s (%s) = %s" %(name, T.__name__, desc))
			
			if act.kv_args:
				synopsis.append("[")
				
				for name, (T, default, desc) in act.kv_args.items():
					synopsis.append("%s=%s" %(name, default or "…"))
					args.append("+ %s (%s) = %s" %(name, T.__name__, desc))
				
				synopsis.append("]")
			
			print("  " + " ".join(synopsis))
			for arg in args:
				print("    " + arg)
			
			print()
			print("      " + act.__doc__.strip())
			
			print()

if __name__ == "__main__":
	class TestActionHello(Action):
		p_args = [("NAME", str, "description"), ("AGE", int, "another short piece of wisdom")]
		kv_args = {"speed": (float, 75, "more stuff")}
		
		def __call__(self, cmd, NAME, AGE, speed):
			print("Hello %s, age %d." %(NAME, AGE))
			print("Called by %s with %s" %(cmd, speed))
	
	routing = {
		"person": {
			"hello": TestActionHello()
		}
	}
	
	invo = Invocation(routing)
	invo.parse_argv(["person", "hello", "Martin", "31", "speed=85"])
	invo.execute()