traffic: Fixed county matching.

1 files changed, 10 insertions, 2 deletions

diff --git a/modules/traffic.py b/modules/traffic.py
index b315660..76e8cc8 100644
--- a/modules/traffic.py
+++ b/modules/traffic.py
@@ -32,6 +32,7 @@ counties = {
 }
 
 class InvalidCountyException(Exception): pass
+class AmbiguousCountyException(Exception): pass
 
 class Module:
 	def __init__(self, bot):
@@ -49,9 +50,14 @@ class Module:
 			rn = None
 			c = args[0]
 		if c:
-			c = c.decode('utf8').capitalize()
-			if not c in counties:
+			c = c.decode('utf8').lower()
+			matches = [s for s in counties if c in s.lower()]
+			if not len(matches):
 				raise InvalidCountyException()
+			elif len(matches) > 1:
+				raise AmbiguousCountyException(matches)
+			else:
+				c = matches[0]
 
 		if rn:
 			if rn.lower().startswith('e'):
@@ -103,6 +109,8 @@ class Module:
 				self.irc.msg(nick.split('!')[0], line.encode('utf8'))
 		except InvalidCountyException:
 			self.irc.msg(channel if not channel == self.irc.nickname else nick.split('!')[0], 'Invalid county.')
+		except AmbiguousCountyException as e:
+			self.irc.msg(channel if not channel == self.irc.nickname else nick.split('!')[0], ('Ambiguous county: %s' % (', '.join(e.message))).encode('utf8'))
 
 if __name__ == '__main__':
 	import sys