From e27f33b2d00839f00054b4be4f7d8b789c559cf9 Mon Sep 17 00:00:00 2001
From: Luca Beltrame <lbeltrame@kde.org>
Date: Sun, 16 Feb 2020 10:10:36 +0100
Subject: [PATCH] Make sure the trip is properly merged

Clean up and make the code more readable.
---
 tripit_add.py | 53 +++++++++++++++++++++++++++++++--------------------
 1 file changed, 32 insertions(+), 21 deletions(-)

diff --git a/tripit_add.py b/tripit_add.py
index 984c44f..21a9fa9 100644
--- a/tripit_add.py
+++ b/tripit_add.py
@@ -78,6 +78,21 @@ def _convert_datetimes(segment: int) -> List[str]:
     return date, time
 
 
+def _parse_record(segment: dict) -> Tuple[str]:
+
+    start_station = segment["start_station"].title()
+    start_date, start_time = _convert_datetimes(segment["start_time"])
+
+    end_station = segment["end_station"].title()
+    end_date, end_time = _convert_datetimes(segment["end_time"])
+    train_type = segment["train_name"]
+    train_number = ("" if segment.get("train_number") is None
+                    else segment["train_number"])
+
+    return (start_station, start_date, start_time, end_station, end_date,
+            end_time, train_type, train_number)
+
+
 def parse_json(json_data: dict, carrier: str = None,
                result_number: int = None) -> str:
 
@@ -89,47 +104,43 @@ def parse_json(json_data: dict, carrier: str = None,
 
     iter_records = mlt.peekable(record["steps"])
 
-    previous_record = defaultdict(None)
+    previous_record = None
 
     for segment in iter_records:
 
-        start_station = segment["start_station"].title()
-        start_date, start_time = _convert_datetimes(segment["start_time"])
-
-        end_station = segment["end_station"].title()
-        end_date, end_time = _convert_datetimes(segment["end_time"])
+        parsed = _parse_record(segment)
 
+        start_station, start_date, start_time, *rest = parsed
+        end_station, end_date, end_time, *rest = rest
+        train_type, train_number = rest
 
         # FIXME: Impossible to extract it from the current Hyperdia data
         carrier = "" if carrier is None else carrier
-        train_type = segment["train_name"]
-        train_number = ("" if segment.get("train_number") is None
-                        else segment["train_number"])
-
         next_record = iter_records.peek(None)
 
         if next_record is not None:
 
+            if segment["go_through"] and next_record["go_through"]:
+                # Nothing to be done here
+                continue
+
             if not segment["go_through"] and next_record["go_through"]:
                 # Next one is going through: set the start,
                 # but don't add anything
-                previous_record["start_station"] = start_station
-                previous_record["start_time"] = start_time
-                previous_record["start_date"] = start_date
-                previous_record["train_name"] = train_type
+                previous_record = segment
                 continue
+
             elif not next_record["go_through"] and segment["go_through"]:
                 # The next one is a "real" one, set the end, and add
-                # the record
-                start_station = previous_record["start_station"]
-                start_time = previous_record["start_time"]
-                start_date = previous_record["start_date"]
-                train_type = previous_record["train_name"]
-                previous_record = dict.fromkeys(previous_record.keys(), None)
+                # use the existing record to extract the data
+                parsed = _parse_record(previous_record)
+                start_station, start_date, start_time, *_ = parsed
+                *_, train_type, train_number = parsed
+                previous_record = None
 
         subdoc = add_train_segment(start_station,
                                    end_station,
-                                   start_time,
+                                start_time,
                                    end_time,
                                    start_date,
                                    train=train_type,