96btool: piechart: Extend the chartable formats

2 files changed, 21 insertions, 10 deletions

diff --git a/bin/96btool b/bin/96btool
index c7489b8..bff847a 100755
--- a/bin/96btool
+++ b/bin/96btool
@@ -509,16 +509,19 @@ def do_passwd(args):
 def do_piechart(args):
 	'''Visualise data as a pie chart'''
 	count = load_json(args.json)
+
+	# Assume the incoming JSON is created by 96btool count
 	try:
-		# Assume the incoming JSON is a dictionary created by
-		# 96btool count
-		chart.piechart(count, args.output, args.title)
-	except AttributeError:
-		# Hmnn... guess this is just a list of posts then...
-		posts = [ Post(p) for p in count ]
-		count = collect.accumulate(posts,
-				lambda p: p['topic']['category'])
 		chart.piechart(count, args.output, args.title)
+		return
+	except KeyError:
+		pass
+
+	# Hmnn... guess this is just a list of posts then...
+	posts = [ Post(p) for p in count ]
+	count = collect.accumulate(posts,
+			lambda p: p['topic']['category'])
+	chart.piechart(count, args.output, args.title)
 
 def do_pull(args):
 	'''Update the local cache of the database'''
diff --git a/lib/python/toys/chart.py b/lib/python/toys/chart.py
index 7f78123..9a02a79 100644
--- a/lib/python/toys/chart.py
+++ b/lib/python/toys/chart.py
@@ -5,6 +5,8 @@ Works best with data sets prepared using toys.collect
 '''
 
 import hashlib
+import matplotlib
+matplotlib.use('Agg')
 import matplotlib.pyplot as plt
 
 import toys.collect as collect
@@ -66,8 +68,14 @@ def stacked_barchart(things, filename, title=None, xlabel=None, ylabel=None):
 	plt.close()
 
 def piechart(things, filename, title=None):
-	labels = sorted(things.keys())
-	counts = [ things[l] for l in labels ]
+	labels = None
+	if isinstance(things, dict):
+		labels = sorted(things.keys())
+		counts = [ things[l] for l in labels ]
+	else:
+		labels = [ t[0] for t in things ]
+		counts = [ t[1] for t in things ]
+
 	colours = [ get_colour(l) for l in labels ]
 
 	wedges, texts = plt.pie(counts, colors=colours, startangle=90)