tracking: Overhauled some code to support other data sources.

Tracking code is now in a separate class, which is chosen by a prefix on
the tracking numbers (prefix:code). By subclassing the TrackingModule
class, new sources can easily be added.

1 files changed, 102 insertions, 47 deletions

diff --git a/modules/tracking.py b/modules/tracking.py
index b00e9fe..4e2a195 100644
--- a/modules/tracking.py
+++ b/modules/tracking.py
@@ -30,12 +30,14 @@ class Consignment(Base):
 	label = Column(String)
 	added = Column(DateTime, nullable = False)
 	last = Column(DateTime)
+	type = Column(String, nullable = False) # Tracking module
 
 	packages = relation('Package', backref = 'consignment', primaryjoin = 'Package.consignment_id == Consignment.id')
 
-	def __init__(self, nick, channel, code, label, added):
+	def __init__(self, nick, channel, type, code, label, added):
 		self.nick = nick
 		self.channel = channel
+		self.type = type
 		self.code = code
 		self.label = label
 		self.added = added
@@ -55,33 +57,26 @@ class Package(Base):
 
 class NoPackageFound(Exception): pass
 
-class Module:
-	url = 'http://sporing.posten.no/sporing.xml?q=%s'
+class TrackingResult:
+	def __init__(self, code, date, desc, delivered = False):
+		self.code, self.date, self.desc, self.delivered = code, date, desc, delivered
 
-	def __init__(self, bot):
-		self.setup_db()
-		self.irc = bot
-		self.tracking = []
-		self.lc = LoopingCall(self.lc_callback)
-		if bot:
-			self.lc.start(config.getfloat(cfg_section, 'interval'), False)
-			self.irc.register_keyword('!track', self)
+	def __str__(self):
+		return '\002{code}\002 {date} - {desc}'.format(code = self.code, date = self.date, desc = self.desc)
 
-	def stop(self):
-		try:
-			self.lc.stop()
-		except:
-			pass
+tracking_modules = {}
 
-	def setup_db(self):
-		global engine, Session
+class TrackingModuleMeta(type):
+	def __init__(cls, name, bases, dict):
+		if name != 'TrackingModule':
+			tracking_modules[cls.name] = cls
 
-		if engine and Session:
-			return
+class TrackingModule(object):
+	__metaclass__ = TrackingModuleMeta
 
-		engine = create_engine(config.get(cfg_section, 'db_path'))
-		Base.metadata.create_all(engine)
-		Session = sessionmaker(bind = engine, autoflush = True, autocommit = False)
+class PostenModule(TrackingModule):
+	name = 'posten'
+	url = 'http://sporing.posten.no/sporing.xml?q=%s'
 
 	def get_xml(self, url):
 		try:
@@ -120,14 +115,59 @@ class Module:
 			if city:
 				desc = '%s (%s)' % (desc, city.encode('utf8'))
 			date = last.find('{%s}OccuredAtDisplayDate' % ns).text + ' ' + last.find('{%s}OccuredAtDisplayTime' % ns).text
-			results.append((code, isodate, desc))
+			results.append(TrackingResult(code, isodate, desc, 'Sendingen er utlevert' in desc))
 		return results
 
+def code_split(code):
+	if ':' in code:
+		return code.split(':')
+	else:
+		return 'posten', code
+
+def get_tracking_module(arg):
+	type, code = code_split(arg)
+
+	return tracking_modules[type]()
+
+class Module:
+	url = 'http://sporing.posten.no/sporing.xml?q=%s'
+
+	def __init__(self, bot):
+		self.setup_db()
+		self.irc = bot
+		self.tracking = []
+		self.lc = LoopingCall(self.lc_callback)
+		if bot:
+			self.lc.start(config.getfloat(cfg_section, 'interval'), False)
+			self.irc.register_keyword('!track', self)
+
+	def stop(self):
+		try:
+			self.lc.stop()
+		except:
+			pass
+
+	def setup_db(self):
+		global engine, Session
+
+		if engine and Session:
+			return
+
+		engine = create_engine(config.get(cfg_section, 'db_path'))
+		Base.metadata.create_all(engine)
+		Session = sessionmaker(bind = engine, autoflush = True, autocommit = False)
+
 	def track_start(self, code, label, nick, channel):
+		type, code = code_split(code)
+		if not type in tracking_modules:
+			return '\002%s\002 is an invalid type.' % type
+		if not len(code):
+			return 'Missing tracking number.'
+
 		msg = None
 		try:
 			session = Session()
-			consignment = Consignment(nick, channel, code, label.decode('utf8'), datetime.datetime.utcnow())
+			consignment = Consignment(nick, channel, type, code, label.decode('utf8'), datetime.datetime.utcnow())
 			session.add(consignment)
 			session.commit()
 			msg = 'Now tracking \002%s\002.' % code
@@ -139,10 +179,16 @@ class Module:
 		return msg
 
 	def track_stop(self, code, nick, channel):
+		type, code = code_split(code)
+		if not len(code):
+			return 'Missing tracking number or label.'
 		msg = None
 		try:
 			session = Session()
-			consignment = session.query(Consignment).filter(or_(Consignment.code == code, Consignment.label == code)).one()
+			consignment = session.query(Consignment).filter(or_(Consignment.code == code, Consignment.label == code))
+			if nick:
+				consignment = consignment.filter(Consignment.nick == nick)
+			consignment = consignment.one()
 			code = consignment.code
 			for p in consignment.packages:
 				session.delete(p)
@@ -157,6 +203,7 @@ class Module:
 		return msg
 
 	def track_status(self, code, nick, channel):
+		# don't code_split() here; this is done later if no results are returned
 		msg = []
 		try:
 			session = Session()
@@ -191,10 +238,12 @@ class Module:
 				msg = results
 			else:
 				try:
-					data = self.track(code)
+					# type:code is handled by get_tracking_module()
+					tm = get_tracking_module(code)
+					data = tm.track(code)
 					if data:
 						for package in data:
-							msg.append('\002%s\002 %s - %s' % (package[0], package[1], package[2].decode('utf8')))
+							msg.append('\002%s\002 %s - %s' % (package.code, package.date, package.desc.decode('utf8')))
 					elif code:
 						msg = 'Failed to fetch tracking data for \002%s\002.' % code.decode('utf8')
 					else:
@@ -206,6 +255,10 @@ class Module:
 		return msg
 
 	def track_label(self, code, label, nick, channel):
+		type, code = code_split(code)
+		if not len(code):
+			return 'Missing tracking number or label.'
+
 		msg = None
 		try:
 			session = Session()
@@ -303,18 +356,19 @@ class Module:
 		if consignment.last and consignment.last > now - td:
 			return
 		consignment.last = now
-		label = (' (%s)' % consignment.label if consignment.label else '')
+		label = (' (%s)' % consignment.label if consignment.label else u'').encode('utf-8')
 		target = consignment.channel or consignment.nick
 		target = target.encode('utf-8')
 		removed = False
-		data_ = self.track(consignment.code)
-		if not data_:
+		tm = tracking_modules[consignment.type]()
+		package_results = tm.track(consignment.code)
+		if not package_results:
 			return
-		for data in data_:
+		for data in package_results:
 			try:
-				package = session.query(Package).filter(and_(Package.consignment_id == consignment.id, Package.code == data[0])).one()
+				package = session.query(Package).filter(and_(Package.consignment_id == consignment.id, Package.code == data.code)).one()
 			except NoResultFound:
-				package = Package(consignment.id, data[0])
+				package = Package(consignment.id, data.code)
 				session.add(package)
 				try:
 					session.commit()
@@ -322,16 +376,16 @@ class Module:
 					# assume several packets within the same consignment was added as a consignment
 					session.rollback()
 					removed = True
-					self.irc.msg(target, '%s: \002%s\002 conflicts with another consignment' % (consignment.nick.encode('utf-8'), data[0]))
+					self.irc.msg(target, '%s: \002%s\002 conflicts with another consignment' % (consignment.nick.encode('utf-8'), data.code))
 					break
-				package = session.query(Package).filter(and_(Package.consignment_id == consignment.id, Package.code == data[0])).one()
+				package = session.query(Package).filter(and_(Package.consignment_id == consignment.id, Package.code == data.code)).one()
 
 			if package.last == None or data[1] > package.last:
-				code = data[0]
-				last = data[1]
-				desc = data[2]
+				code = data.code
+				last = data.date
+				desc = data.desc
 				msg = '%s: \002%s\002%s %s - %s' % (consignment.nick.encode('utf-8'), package.code.encode('utf-8'), label.encode('utf-8'), last, desc)
-				if desc.startswith('Sendingen er utlevert'):
+				if data.delivered:
 					session.delete(package)
 					msg += ' (Package delivered - tracking stopped)'
 					removed = True
@@ -342,18 +396,18 @@ class Module:
 				if announce:
 					self.irc.msg(target, msg)
 			elif package.last != None and package.last < datetime.datetime.utcnow() - datetime.timedelta(hours = 24*30):
-				msg = '%s: Removing stale package \002%s\002%s' % (consignment.nick, consignment.code, label)
-				self.irc.msg(target, msg.encode('utf-8'))
+				msg = '%s: Removing stale package \002%s\002%s' % (consignment.nick.encode('utf-8'), consignment.code.encode('utf-8'), label)
+				self.irc.msg(target, msg)
 				session.delete(package)
 				removed = True
 		# empty consignment and 30 days old
 		if len(consignment.packages) == 0 and consignment.added < datetime.datetime.utcnow() - datetime.timedelta(hours = 24*30):
-			msg = '%s: Removing stale consignment \002%s\002%s' % (consignment.nick, consignment.code, label)
-			self.irc.msg(target, msg.encode('utf-8'))
+			msg = '%s: Removing stale consignment \002%s\002%s' % (consignment.nick.encode('utf-8'), consignment.code.encode('utf-8'), label)
+			self.irc.msg(target, msg)
 			session.delete(consignment)
 		if removed and len(consignment.packages) == 0:
-			msg = '%s: \002%s\002%s is no longer being tracked' % (consignment.nick, consignment.code, label)
-			self.irc.msg(target, msg.encode('utf-8'))
+			msg = '%s: \002%s\002%s is no longer being tracked' % (consignment.nick.encode('utf-8'), consignment.code.encode('utf-8'), label)
+			self.irc.msg(target, msg)
 			session.delete(consignment)
 
 	def lc_callback(self):
@@ -372,7 +426,8 @@ if __name__ == '__main__':
 	config.read([os.path.expanduser('~/.fot')])
 	m = Module(None)
 	for code in sys.argv[1:]:
-		tracking = m.track(code)
+		tm = get_tracking_module(code)
+		tracking = tm.track(code)
 		if tracking:
 			for track in tracking:
 				if type(track) == list: