From 772d475015f99a118bc86f5ac267a7ace264bda3 Mon Sep 17 00:00:00 2001
From: Christopher Teutsch <chris@heychris.eu>
Date: Sat, 22 Jun 2019 01:36:35 +0200
Subject: [PATCH] Write line names to separate table

---
 .gitignore |  2 ++
 monitor.py | 55 +++++++++++++++++++++++++++++++++++++++++++++---------
 2 files changed, 48 insertions(+), 9 deletions(-)
 create mode 100644 .gitignore

diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..4abbfe9
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1,2 @@
+cache/
+*.pyc
\ No newline at end of file
diff --git a/monitor.py b/monitor.py
index d41a0c2..bb3ed57 100755
--- a/monitor.py
+++ b/monitor.py
@@ -1,7 +1,8 @@
 #!/usr/bin/env python3
 import requests
 import json
-from typing import List
+from typing import List, Dict
+from io import BytesIO
 # from pprint import pprint, pformat
 import datetime
 import pause
@@ -9,8 +10,10 @@ import sys
 import mysql.connector
 import configparser
 import argparse
+from cachecontrol import CacheControl
+from cachecontrol.caches.file_cache import FileCache
 
-TABLE = """
+VRR_TABLE = """
 CREATE TABLE IF NOT EXISTS vrr (
 line_code varchar(9) not null,
 direction_code varchar(1) not null,
@@ -20,6 +23,11 @@ status enum('on_time', 'early', 'late', 'cancelled', 'no_data') not null,
 delay_value int,
 primary key (line_code, direction_code, station_id, orig_datetime));
 """
+VRR_LINES_TABLE = """
+CREATE TABLE IF NOT EXISTS vrr_lines (
+line_code varchar(9) not null primary key,
+line_number varchar(6) not null);
+"""
 
 
 class MOT:
@@ -33,8 +41,16 @@ class MOT:
     ALL_MODES = [LONG_DISTANCE_TRAIN, REGIONAL_TRAIN, COMMUTER_TRAIN, UNDERGROUND_TRAIN, TRAM, BUS, ELEVATED_TRAIN]
 
 
-last_reply = requests.Response
+class Line:
+    line_code = str
+    line_name = str
 
+    def __init__(self, code, name):
+        self.line_code = code
+        self.line_name = name
+
+
+last_reply = None
 parser = argparse.ArgumentParser()
 parser.add_argument("-s", "--setup", help="Run the setup routine.", action="store_true", dest="setup")
 
@@ -88,7 +104,7 @@ TRIP_CANCELLED = -9999
 
 # Initialize Requests session
 
-HTTP = requests.session()
+HTTP = CacheControl(requests.session(), cache=FileCache('.cache'))
 
 
 def make_request_data(station_id: int, result_count: int = 8, modes: List = MOT.ALL_MODES,
@@ -185,6 +201,19 @@ def fixup_data(d: dict) -> dict:
     return d
 
 
+def make_linecode_table(data: dict) -> Dict[str, str]:
+    codes = {}
+
+    for trip in data['departureData']:
+        if trip['lineCode'] in codes:
+            if trip['lineNumber'] != codes[trip['lineCode']]:
+                print('lineCode {} already present as {}! Replacing it with {}.'.format(
+                    trip['lineCode'], codes[trip['lineCode']], trip['lineNumber'])
+                )
+        codes[trip['lineCode']] = trip['lineNumber']
+    return codes
+
+
 def print_trip(trip: dict, full_text: bool = False) -> None:
     if full_text:
         fmt = {
@@ -281,11 +310,13 @@ def update(station_id: int, lines: List[str]) -> dict or None:
     return reply_data
 
 
-def wait(cxn: mysql.connector.MySQLConnection, station_id: int):
+def wait(station_id: int):
+    cxn = mysql.connector.MySQLConnection(**db_config)
     cur = cxn.cursor()
     while True:
         data = update(station_id, USE_LINES)
         if data is not None:
+            line_codes = make_linecode_table(data)
             for t in data['departureData']:
                 cur.execute('REPLACE INTO vrr '
                             '(line_code, direction_code, station_id, orig_datetime, status, delay_value)'
@@ -293,7 +324,11 @@ def wait(cxn: mysql.connector.MySQLConnection, station_id: int):
                             'from_unixtime(%s), %s, %s)',
                             (t['lineCode'], t['directionCode'], station_id,
                              t['orgFullTime'], _make_status_value_for_sql(t), _make_delay_value_for_sql(t['delay'])))
-            cxn.commit()
+            for k, v in line_codes.items():
+                cur.execute('REPLACE INTO vrr_lines'
+                            '(line_code, line_number)'
+                            'VALUES (%s, %s)', (k, v))
+        cxn.commit()
         next_refresh = get_next_refresh(data)
         print("Sleeping until " + datetime.datetime.fromtimestamp(next_refresh).isoformat(), file=sys.stderr)
         pause.until(next_refresh)
@@ -302,10 +337,12 @@ def wait(cxn: mysql.connector.MySQLConnection, station_id: int):
 def main():
     cxn = mysql.connector.connect(**db_config)
     cursor = cxn.cursor()
-    cursor.execute(TABLE)
-
+    cursor.execute(VRR_TABLE)
+    cursor.execute(VRR_LINES_TABLE)
+    cursor.close()
+    cxn.close()
     try:
-        wait(cxn, USE_STATION_ID)
+        wait(USE_STATION_ID)
     except (ValueError, TypeError) as e:
         with open("fault.json", "wb") as o:
             if last_reply is not None: