tracking: stop and status now handles labels.

1 files changed, 18 insertions, 21 deletions

diff --git a/modules/tracking.py b/modules/tracking.py
index 4bd58be..405514e 100644
--- a/modules/tracking.py
+++ b/modules/tracking.py
@@ -14,6 +14,7 @@ from sqlalchemy.ext.declarative import declarative_base
 from sqlalchemy.orm import sessionmaker, relation, backref
 from sqlalchemy.orm.exc import NoResultFound
 from sqlalchemy.exc import IntegrityError
+from sqlalchemy.sql import or_
 
 engine = Session = None
 Base = declarative_base()
@@ -136,15 +137,16 @@ class Module:
 		msg = None
 		try:
 			session = Session()
-			consignment = session.query(Consignment).filter_by(code = code).one()
+			consignment = session.query(Consignment).filter(or_(Consignment.code == code, Consignment.label == code)).one()
+			code = consignment.code
 			for p in consignment.packages:
 				session.delete(p)
 			label = (' (%s)' % consignment.label) if consignment.label else ''
 			session.delete(consignment)
 			session.commit()
-			msg = 'No longer tracking \002%s\002%s' % (code, label)
+			msg = 'No longer tracking \002%s\002%s' % (code.decode('utf8'), label)
 		except NoResultFound:
-			msg = '\002%s\002 is not being tracked.' % code
+			msg = '\002%s\002 is not being tracked.' % code.decode('utf8')
 		finally:
 			session.close()
 		return msg
@@ -153,7 +155,7 @@ class Module:
 		msg = []
 		try:
 			session = Session()
-			consignments = session.query(Consignment).filter(Consignment.code.like(code + '%'))
+			consignments = session.query(Consignment).filter(or_(Consignment.code.like(code + '%'), Consignment.label.like((code + '%'))))
 			if nick:
 				consignments = consignments.filter(Consignment.nick == nick)
 			results = []
@@ -188,7 +190,7 @@ class Module:
 						for package in data:
 							msg.append('\002%s\002 %s - %s' % (package[0], package[1], package[2].decode('utf8')))
 					elif code:
-						msg = 'Failed to fetch tracking data for \002%s\002.' % code
+						msg = 'Failed to fetch tracking data for \002%s\002.' % code.decode('utf8')
 					else:
 						msg = 'No tracking number given or registered.'
 				except Exception as e:
@@ -232,30 +234,25 @@ class Module:
 		return msg
 
 	def keyword(self, nick, channel, kw, msg):
+		usage = 'Usage: !track (start|label TRACKINGNO [LABEL])|(stop|status TRACKINGNO|LABEL)|list'
 		args = msg.split()
 		if len(args) < 1:
-			self.irc.msg(channel if not channel == self.irc.nickname else nick.split('!')[0],
-					'Usage: !track (start|stop|status|label TRACKINGNO [LABEL])|list')
+			self.irc.msg(channel if not channel == self.irc.nickname else nick.split('!')[0], usage)
 			return
 
 		mode = args[0]
-		if mode.lower() == 'start':
-			if len(args) < 2:
-				self.irc.msg(channel if not channel == self.irc.nickname else nick.split('!')[0],
-						'Usage: !track (start|stop|status|label TRACKINGNO [LABEL])|list')
+		if mode.lower() in ('start', 'stop', 'status'):
+			if len(args) < 2 and mode.lower() != 'status':
+				self.irc.msg(channel if not channel == self.irc.nickname else nick.split('!')[0], usage)
 				return
-			code = args[1]
-			label = ' '.join(args[2:])
-		elif mode.lower() in ('stop', 'status'):
-			if len(args) != 2 and mode.lower() != 'status':
-				self.irc.msg(channel if not channel == self.irc.nickname else nick.split('!')[0],
-						'Usage: !track (start|stop|status|label TRACKINGNO [LABEL])|list')
-				return
-			code = args[1] if len(args) == 2 else ''
+			if mode.lower() == 'start':
+				code = args[1]
+				label = ' '.join(args[2:]) if len(args) > 2 else ''
+			else:
+				code = ' '.join(args[1:])
 		elif mode.lower() == 'label':
 			if len(args) < 3:
-				self.irc.msg(channel if not channel == self.irc.nickname else nick.split('!')[0],
-						'Usage: !track (start|stop|status|label TRACKINGNO [LABEL])|list')
+				self.irc.msg(channel if not channel == self.irc.nickname else nick.split('!')[0], usage)
 				return
 			code = args[1]
 			label = ' '.join(args[2:])