diff options
| -rw-r--r-- | modules/tracking.py | 149 |
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: |