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