From 9190d399846bddfdf093f188cc2292a0ca31cc08 Mon Sep 17 00:00:00 2001
From: Marco Aceti <mail@marcoaceti.it>
Date: Thu, 23 Mar 2023 14:50:51 +0100
Subject: [PATCH] Add geojson station data extractor

---
 main.py                  |  4 +---
 requirements.txt         |  1 +
 src/station_extractor.py | 33 +++++++++++++++++++++++++++++++++
 3 files changed, 35 insertions(+), 3 deletions(-)

diff --git a/main.py b/main.py
index fe990f7..3eac3f1 100644
--- a/main.py
+++ b/main.py
@@ -65,9 +65,7 @@ s_extractor.add_argument(
 s_extractor.add_argument(
     "-f",
     default="csv",
-    choices=[
-        "csv",
-    ],
+    choices=["csv", "geojson"],
     help="output file format",
     dest="format",
 )
diff --git a/requirements.txt b/requirements.txt
index 63d484f..2760fff 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,3 +1,4 @@
 requests==2.28.2
 python-dateutil==2.8.2
 tqdm==4.65.0
+geojson==3.0.1
diff --git a/src/station_extractor.py b/src/station_extractor.py
index bbceaa0..8b04689 100644
--- a/src/station_extractor.py
+++ b/src/station_extractor.py
@@ -3,6 +3,8 @@ import csv
 import pickle
 from pathlib import Path
 
+from geojson import Feature, FeatureCollection, Point
+
 from src.scraper.station import Station
 from src.utils import parse_input_format_output_args
 
@@ -59,11 +61,42 @@ def to_csv(data: dict[str, Station], output_file: Path) -> None:
                 station.position[1] if station.position else None,
             )
         )
+    csvfile.close()
+
+
+def to_geojson(data: dict[str, Station], output_file: Path) -> None:
+    feature_list: list[Feature] = list()
+
+    for station_c in data:
+        station: Station = data[station_c]
+        if not station.position:
+            continue
+
+        feature: Feature = Feature(
+            geometry=Point((station.position[1], station.position[0])),
+            properties={
+                "code": station.code,
+                "name": station.name,
+                "short_name": station.short_name
+                if hasattr(station, "short_name")
+                else None,
+                "region": station.region_code,
+            },
+        )
+        feature_list.append(feature)
+
+    collection: FeatureCollection = FeatureCollection(feature_list)
+    with open(output_file, "w+") as f:
+        f.write(str(collection))
 
 
 def main(args: argparse.Namespace):
     input_f, output_f, format = parse_input_format_output_args(args)
 
     data: dict[str, Station] = load_file(input_f)
+
     if format == "csv":
         to_csv(data, output_f)
+
+    if format == "geojson":
+        to_geojson(data, output_f)
-- 
GitLab