tracking: Implemented label support for consignments.

1 files changed, 74 insertions, 14 deletions

diff --git a/modules/tracking.py b/modules/tracking.py
index d2560cd..aa57763 100644
--- a/modules/tracking.py
+++ b/modules/tracking.py
@@ -25,14 +25,16 @@ class Consignment(Base):
 	nick = Column(String, nullable = False)
 	channel = Column(String) # NULL for private
 	code = Column(String, nullable = False, unique = True)
+	label = Column(String)
 	added = Column(DateTime, nullable = False)
 
 	packages = relation('Package', backref = 'consignment', primaryjoin = 'Package.consignment_id == Consignment.id')
 
-	def __init__(self, nick, channel, code, added):
+	def __init__(self, nick, channel, code, label, added):
 		self.nick = nick
 		self.channel = channel
 		self.code = code
+		self.label = label
 		self.added = added
 
 class Package(Base):
@@ -116,11 +118,11 @@ class Module:
 			results.append((code, isodate, desc))
 		return results
 
-	def track_start(self, code, nick, channel):
+	def track_start(self, code, label, nick, channel):
 		msg = None
 		try:
 			session = Session()
-			consignment = Consignment(nick, channel, code, datetime.datetime.utcnow())
+			consignment = Consignment(nick, channel, code, label.decode('utf8'), datetime.datetime.utcnow())
 			session.add(consignment)
 			session.commit()
 			msg = 'Now tracking \002%s\002.' % code
@@ -137,9 +139,10 @@ class Module:
 			consignment = session.query(Consignment).filter_by(code = code).one()
 			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' % code
+			msg = 'No longer tracking \002%s\002%s' % (code, label)
 		except NoResultFound:
 			msg = '\002%s\002 is not being tracked.' % code
 		finally:
@@ -156,17 +159,39 @@ class Module:
 			results = []
 			for row in consignments:
 				i = 0
+				if row.label:
+					label_header = True
+					results.append('Packages for \002%s\002:' % row.label)
+				else:
+					label_header = False
 				for package in row.packages:
 					i += 1
 					date = package.last
 					desc = package.status
+
+					if package.code == row.code and row.label:
+						label = ' (%s)' % row.label
+					else:
+						label = ''
+
 					if date and desc:
-						s = '\002%s\002 %s - %s' % (package.code, date, desc)
+						s = '\002%s\002%s %s - %s' % (package.code, label, date, desc)
 					else:
-						s = 'No tracking info found for \002%s\002' % package.code
+						s = 'No tracking info found for \002%s\002%s' % (package.code, label)
+
 					results.append(s)
+
 				if i == 0:
-					results.append('No packages found for \002%s\002' % row.code)
+					if row.label:
+						label = ' (%s)' % row.label
+					else:
+						label = ''
+					results.append('No packages found for \002%s\002%s' % (row.code, label))
+
+				if label_header:
+					del results[0]
+					label_header = False
+
 			if len(results):
 				msg = results
 			else:
@@ -185,6 +210,19 @@ class Module:
 			session.close()
 		return msg
 
+	def track_label(self, code, label, nick, channel):
+		msg = None
+		try:
+			session = Session()
+			consignment = session.query(Consignment).filter_by(code = code).filter_by(nick = nick).one()
+			consignment.label = label.decode('utf8')
+			session.add(consignment)
+			session.commit()
+			msg = '\002%s\002 now labelled \002%s\002' % (code, label.decode('utf8'))
+		finally:
+			session.close()
+		return msg
+
 	def track_list(self, nick):
 		msg = None
 		try:
@@ -192,7 +230,10 @@ class Module:
 			trackings = []
 			consignments = session.query(Consignment).filter_by(nick = nick)
 			for row in consignments:
-				trackings.append(row.code)
+				s = row.code
+				if row.label:
+					s += ' (%s)' % row.label
+				trackings.append(s)
 			if len(trackings):
 				msg = 'Trackings for %s: %s' % (nick, ', '.join(trackings))
 			else:
@@ -206,25 +247,43 @@ class Module:
 	def keyword(self, nick, channel, kw, msg):
 		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 TRACKINGNO)|list')
+			self.irc.msg(channel if not channel == self.irc.nickname else nick.split('!')[0],
+					'Usage: !track (start|stop|status|label TRACKINGNO [LABEL])|list')
 			return
 
 		mode = args[0]
-		if mode.lower() in ('start', 'stop', 'status'):
+		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')
+				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 TRACKINGNO)|list')
+				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 ''
+		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')
+				return
+			code = args[1]
+			label = ' '.join(args[2:])
 
 		msg = None
 		if mode.lower() == 'start':
-			msg = self.track_start(code, nick.split('!')[0], channel if not channel == self.irc.nickname else None)
+			msg = self.track_start(code, label, nick.split('!')[0], channel if not channel == self.irc.nickname else None)
 		elif mode.lower() == 'stop':
 			msg = self.track_stop(code, nick.split('!')[0], channel if not channel == self.irc.nickname else None)
 		elif mode.lower() == 'status':
 			msg = self.track_status(code, nick.split('!')[0], channel if not channel == self.irc.nickname else None)
 		elif mode.lower() == 'list':
 			msg = self.track_list(nick.split('!')[0])
+		elif mode.lower() == 'label':
+			msg = self.track_label(code, label, nick.split('!')[0], channel if not channel == self.irc.nickname else None)
 		else:
 			msg = 'Invalid mode "%s".' % mode
 
@@ -260,7 +319,8 @@ class Module:
 						code = data[0]
 						last = data[1]
 						desc = data[2]
-						msg = '%s: \002%s\002 %s - %s' % (row.nick.encode('utf-8'), package.code.encode('utf-8'), last, desc)
+						label = (' (%s)' % row.label.encode('utf-8') if row.label else '')
+						msg = '%s: \002%s\002 %s - %s' % (row.nick.encode('utf-8'), package.code.encode('utf-8'), label, last, desc)
 						if desc.startswith('Sendingen er utlevert'):
 							session.delete(package)
 							msg += ' (Package delivered - tracking stopped)'
@@ -271,7 +331,7 @@ class Module:
 							session.add(package)
 						self.irc.msg(target, msg)
 				if removed and len(row.packages) == 0:
-					msg = '%s: \002%s\002 is no longer being tracked' % (row.nick, row.code)
+					msg = '%s: \002%s\002%s is no longer being tracked' % (row.nick, row.code, label)
 					self.irc.msg(target, msg.encode('utf-8'))
 					session.delete(row)
 			session.commit()