oebb_py/classes/connection.py

import datetime


class Connection:
    def __init__(self, details=False):
        self.services = []
        self.details = details
        self.via = []

    def addService(self, service):
        self.services.append(service)

    def addVia(self, station):
        self.via.append(station)

    def depst(self):
        try:
            return self.services[0].depst
        except:
            return None

    def arrst(self):
        try:
            return self.services[-1].arrst
        except:
            return None

    def deptime(self):
        try:
            return self.services[0].deptime
        except:
            return None

    def arrtime(self):
        try:
            return self.services[-1].arrtime
        except:
            return None

    def currdep(self):
        try:
            return self.services[0].currdep
        except:
            return None

    def currarr(self):
        try:
            return self.services[-1].currarr
        except:
            return None

    def duration(self):
        try:
            return self.services[-1].arrtime - self.services[0].deptime
        except:
            return None

    def durationString(self):
        hrs, scs = divmod(self.duration().total_seconds(), 3600)
        mns, rmd = divmod(scs, 60)

        return "%i:%s" % (hrs, str(int(mns)).zfill(2))

    def changes(self):
        return max(len([service for service in self.services if service.name != "Walk"]) - 1, 0) if self.details else max(len([service for service in self.services[0].name.split("/") if service != "Walk"]) - 1, 0)

    def xml(self, indent=0, cid=False, frm=True, to=True, deptime=True, arrtime=True, duration=True, changes=True, services=True, via=True, servicekwargs={}, stationkwargs={}):
        out = " " * indent + \
            "<connection%s>\n" % ("" if cid is False else " id=\"%i\"" % cid)

        if frm and self.depst():
            out += " " * indent + " <from>\n"
            out += self.depst().xml(indent + 2, **stationkwargs) + "\n"
            out += " " * indent + " </from>\n"

        if via and self.via:
            out += " " * indent + " <via>\n"

            for vst in self.via:
                out += vst.xml(indent + 2, **stationkwargs) + "\n"

            out += " " * indent + " </via>\n"

        if to and self.arrst():
            out += " " * indent + " <to>\n"
            out += self.arrst().xml(indent + 2, **stationkwargs) + "\n"
            out += " " * indent + " </to>\n"

        if (deptime and self.deptime()) or (arrtime and self.arrtime()) or (duration and self.duration()) or (changes and self.changes()):
            out += " " * indent + " <details>\n"

            if deptime and self.deptime():
                out += " " * indent + "  <departure>\n"
                out += " " * indent + \
                    "   <date>%s</date>\n" % datetime.datetime.strftime(
                        self.deptime(), "%d.%m.%Y")
                out += " " * indent + \
                    "   <time>%s</time>\n" % datetime.datetime.strftime(
                        self.deptime(), "%H:%M")
                out += " " * indent + "  </departure>\n"

            if arrtime and self.arrtime():
                out += " " * indent + "  <arrival>\n"
                out += " " * indent + \
                    "   <date>%s</date>\n" % datetime.datetime.strftime(
                        self.arrtime(), "%d.%m.%Y")
                out += " " * indent + \
                    "   <time>%s</time>\n" % datetime.datetime.strftime(
                        self.arrtime(), "%H:%M")
                out += " " * indent + "  </arrival>\n"

            out += (" " * indent + "  <duration>%s</duration>\n" %
                    self.durationString()) if duration and self.duration() else ""
            out += (" " * indent + "  <changes>%i</changes>\n" %
                    self.changes()) if changes and self.changes() else ""

            out += " " * indent + " </details>\n"

        if services and self.services:
            out += " " * indent + " <services>\n"

            for i in range(len(self.services)):
                out += self.services[i].xml(indent + 2,
                                            i, **servicekwargs) + "\n"

            out += " " * indent + " </services>\n"

        out += " " * indent + "</connection>"

        return out

    def json(self, indent=0, cid=False, frm=True, to=True, deptime=True, arrtime=True, duration=True, changes=True, services=True, via=True, servicekwargs={}, stationkwargs={}):
        out = " " * indent + "{\n"

        out += (" " * indent + " \"@id\": %i,\n" %
                cid) if cid is not False else ""

        if frm and self.depst():
            out += " " * indent + " \"from\":\n"
            out += self.depst().json(indent + 2, **stationkwargs) + ",\n"

        if via and self.via:
            out += " " * indent + " \"via\": [\n"

            for vst in self.via:
                out += vst.json(indent + 3, **stationkwargs) + ",\n"

            out = "".join(out.rsplit(",", 1))
            out += " " * indent + " ],\n"

        if to and self.arrst():
            out += " " * indent + " \"to\":\n"
            out += self.arrst().json(indent + 2, **stationkwargs) + ",\n"

        if (deptime and self.deptime()) or (arrtime and self.arrtime()) or (duration and self.duration()) or (changes and self.changes()):
            det = ""
            det += " " * indent + " \"details\": {\n"

            if deptime and self.deptime():
                det += " " * indent + "  \"departure\": {\n"
                det += " " * indent + \
                    "   \"date\": \"%s\",\n" % datetime.datetime.strftime(
                        self.deptime(), "%d.%m.%Y")
                det += " " * indent + \
                    "   \"time\": \"%s\"\n" % datetime.datetime.strftime(
                        self.deptime(), "%H:%M")
                det += " " * indent + "  },\n"

            if arrtime and self.arrtime():
                det += " " * indent + "  \"arrival\": {\n"
                det += " " * indent + \
                    "   \"date\": \"%s\",\n" % datetime.datetime.strftime(
                        self.arrtime(), "%d.%m.%Y")
                det += " " * indent + \
                    "   \"time\": \"%s\"\n" % datetime.datetime.strftime(
                        self.arrtime(), "%H:%M")
                det += " " * indent + "  },\n"

            det += (" " * indent + "  \"duration\": \"%s\",\n" %
                    self.durationString()) if duration and self.duration() else ""
            det += (" " * indent + "  \"changes\": %i,\n" %
                    self.changes()) if changes and self.changes() else ""

            det = "".join(det.rsplit(",", 1))
            out += det

            out += " " * indent + " },\n"

        if services and self.services:
            out += " " * indent + " \"services\": [\n"

            for i in range(len(self.services)):
                out += self.services[i].json(indent + 2, i, **servicekwargs) + (
                    ",\n" if not i == len(self.services) - 1 else "\n")

            out += " " * indent + " ],\n"

        out += " " * indent + "}"

        out = "".join(out.rsplit(",", 1))

        return out