pytest/plib/station.py

from typing import List, Union
from plib import Point


class StationError(Exception):
    pass


class Station:
    def __init__(self, code: Union[str, None] = None,
                 height: Union[float, None] = None,
                 location: Union[List[float], None] = None,
                 status: Union[str, None] = None,
                 xyz: Union[List[float], None] = None):
        if code is None or height is None or location is None or status is None or xyz is None:
            raise StationError
        self.code = code
        self.height = height
        self.location = Point(location['lat'], location['lon'])
        self.status = status
        self.xyz = xyz

    def distance_to(self, other: "Station") -> float:
        return self.location.distance_to(other.location)


class StationMap:
    def __init__(self, stations: Union[List[Station], None] = None):
        if stations is None:
            raise StationError
        self.stations = stations

    @classmethod
    def from_json(cls, text):
        stations = []
        for st in text:
            stations.append(dict(st))
        return cls(stations)

    def to_json(self, stations):
        return stations