From 20092be17b574bf2f9374d6129673958a38a9f5b Mon Sep 17 00:00:00 2001
From: n1nj4sec <contact@n1nj4.eu>
Date: Sun, 11 Oct 2015 16:48:00 +0200
Subject: [PATCH] automatically add list completer to argparse with a "choices"
 kw

---
 pupy/pupylib/PupyCompleter.py | 7 +++----
 pupy/pupylib/PupyModule.py    | 4 +++-
 2 files changed, 6 insertions(+), 5 deletions(-)

diff --git a/pupy/pupylib/PupyCompleter.py b/pupy/pupylib/PupyCompleter.py
index 52b4f5d2..22fa99da 100644
--- a/pupy/pupylib/PupyCompleter.py
+++ b/pupy/pupylib/PupyCompleter.py
@@ -26,7 +26,8 @@ def debug(msg):
 
 def list_completer(l):
 	def func(text, line, begidx, endidx):
-		return [x for x in l if text.startswith(x)]
+		return [x+" " for x in l if x.startswith(text)]
+	return func
 
 def void_completer(text, line, begidx, endidx):
 	return []
@@ -100,7 +101,6 @@ class PupyModCompleter(object):
 	def __init__(self):
 		self.conf= {
 			"positional_args":[
-				#TODO
 			],
 			"optional_args":[
 			],
@@ -140,7 +140,6 @@ class PupyModCompleter(object):
 			return None
 
 	def get_positional_arg_index(self, text, line, begidx, endidx):
-		#TODO
 		tab=shlex.split(line)
 		positional_index=-1
 		for i in range(0, len(tab)):
@@ -157,6 +156,7 @@ class PupyModCompleter(object):
 
 	def get_optional_args_completer(self, name):
 		return [x[1]["completer"] for x in self.conf["optional_args"] if x[0]==name][0]
+
 	def get_positional_args_completer(self, index):
 		return self.conf["positional_args"][index][1]["completer"]
 
@@ -171,7 +171,6 @@ class PupyModCompleter(object):
 			return [x+" " for x in self.get_optional_args() if x.startswith(text)]
 		else:
 			try:
-				debug("here")
 				positional_index=self.get_positional_arg_index(text, line, begidx, endidx)-2 # -2 for "run" + "module_name"
 				debug("positional index is %s"%positional_index)
 				return self.get_positional_args_completer(positional_index)(text, line, begidx, endidx)
diff --git a/pupy/pupylib/PupyModule.py b/pupy/pupylib/PupyModule.py
index 4d863861..724d5adc 100644
--- a/pupy/pupylib/PupyModule.py
+++ b/pupy/pupylib/PupyModule.py
@@ -16,7 +16,7 @@
 import argparse
 import sys
 from .PupyErrors import PupyModuleExit
-from .PupyCompleter import PupyModCompleter, void_completer
+from .PupyCompleter import PupyModCompleter, void_completer, list_completer
 import StringIO
 
 class PupyArgumentParser(argparse.ArgumentParser):
@@ -30,6 +30,8 @@ class PupyArgumentParser(argparse.ArgumentParser):
 		if "completer" in kwargs:
 			completer_func=kwargs["completer"]
 			del kwargs["completer"]
+		elif "choices" in kwargs:
+			completer_func=list_completer(kwargs["choices"])
 		else:
 			completer_func=void_completer
 		argparse.ArgumentParser.add_argument(self, *args, **kwargs)