Merge branch 'master' of https://bitbucket.org/Halogenic/wikiquery

1 files changed, 23 insertions, 7 deletions

diff --git a/wikiquery b/wikiquery
index 7682d78..10356dd 100755
--- a/wikiquery
+++ b/wikiquery
@@ -8,17 +8,19 @@ from functools import partial
 
 version = "0.0.2"
 
-def display(o):
+def stringify(o):
 	if isinstance(o, list):
-
 		output = []
 		for item in o:
 			output.append(item.encode("utf-8", "replace"))
 
-		print "'%s'" % "', '".join(output)
+		return '"%s"' % '", "'.join(output)
 	else:
-		print o.encode("utf-8", "replace")
+		return o.encode("utf-8", "replace")
 
+def display(o):
+	print stringify(o)
+	
 def parse_args():
 	parser = optparse.OptionParser(usage="!wiki <topic> [--search|--url]")
 
@@ -36,13 +38,21 @@ def main():
 		sys.exit(0)
 	
 	if args:
-		topic = args[0]
+		topic = args.pop(0)
+		index = int(args.pop(0)) if args else None
+
 		if options.search:
 			results = wikipedia.search(topic)
 			if len(results) == 1:
 				display(wikipedia.summary(results[0]))
 			else:
-				display(results)
+				if index is not None:
+					try:
+						display(wikipedia.summary(results[index]))
+					except IndexError:
+						print "index out of range, options are: [%s]" % stringify(results)
+				else:
+					display(results)
 		elif options.url:
 			page = wikipedia.page(topic)
 			display(page.url)
@@ -50,7 +60,13 @@ def main():
 			try:
 				display(wikipedia.summary(topic))
 			except wikipedia.exceptions.DisambiguationError as e:
-				display(e.options)
+				if index is not None:
+					try:
+						display(wikipedia.summary(e.options[index]))
+					except IndexError:
+						print "index out of range, options are: [%s]" % stringify(e.options)
+				else:
+					display(e.options)
 
 if __name__ == "__main__":
 	main()