Add type annotations for Airport, City

duffel_api/models/__init__.py

@@ -1,5 +1,6 @@
 from .aircraft import Aircraft
-from .airport import Airport, City
+from .airport import Airport
+from .city import City
 from .airline import Airline
 from .loyalty_programme_account import LoyaltyProgrammeAccount
 from .passenger import Passenger

duffel_api/models/aircraft.py

@@ -1,6 +1,19 @@
+from dataclasses import dataclass
+
+
+@dataclass
 class Aircraft:
     """Aircraft are used to describe what passengers will fly in for a given trip"""
 
-    def __init__(self, json):
-        for key in json:
-            setattr(self, key, json[key])
+    id: str
+    iata_code: str
+    name: str
+
+    @classmethod
+    def from_json(cls, json: dict):
+        """Construct a class instance from a JSON response."""
+        return cls(
+            id=json["id"],
+            iata_code=json["iata_code"],
+            name=json["name"],
+        )

duffel_api/models/airport.py

@@ -1,21 +1,34 @@
-class Airport:
-    """Airports are used to identify origins and destinations in journey slices"""
+from dataclasses import dataclass
+from typing import Optional
+
+from ..models import City
 
-    def __init__(self, json):
-        for key in json:
-            value = json[key]
-            if key == "city" and value:
-                setattr(self, key, City(value))
-            else:
-                setattr(self, key, value)
 
+@dataclass
+class Airport:
+    """Airports are used to identify origins and destinations in journey slices"""
 
-class City:
-    """The metropolitan area where the airport is located.
-    Only present for airports which are registered with IATA as
-    belonging to a metropolitan area.
-    """
+    id: str
+    name: str
+    iata_code: Optional[str]
+    icao_code: Optional[str]
+    country_code: str
+    latitude: float
+    longitude: float
+    time_zone: str
+    city: City
 
-    def __init__(self, json):
-        for key in json:
-            setattr(self, key, json[key])
+    @classmethod
+    def from_json(cls, json: dict):
+        """Construct a class instance from a JSON response."""
+        return cls(
+            id=json["id"],
+            name=json["name"],
+            iata_code=json.get("iata_code"),
+            icao_code=json.get("icao_code"),
+            country_code=json["country_code"],
+            latitude=json["latitude"],
+            longitude=json["longitude"],
+            time_zone=json["time_zone"],
+            city=City.from_json(json=json["name"]),
+        )

duffel_api/models/city.py

@@ -0,0 +1,24 @@
+from dataclasses import dataclass
+
+
+@dataclass
+class City:
+    """The metropolitan area where the airport is located.
+    Only present for airports which are registered with IATA as
+    belonging to a metropolitan area.
+    """
+
+    id: str
+    name: str
+    iata_code: str
+    iata_country_code: str
+
+    @classmethod
+    def from_json(cls, json: dict):
+        """Construct a class instance from a JSON response."""
+        return cls(
+            id=json["id"],
+            name=json["name"],
+            iata_code=json["iata_code"],
+            iata_country_code=json["iata_country_code"],
+        )

duffel_api/models/place.py

@@ -1,32 +1,67 @@
-from ..models import City
+from dataclasses import dataclass
+from typing import Optional, Sequence
 
+from ..models import Airport, City
+from ..utils import get_and_transform
 
+
+class CityAirport:
+    """The airport associated to a city."""
+
+    def __init__(self, json):
+        for key in json:
+            value = json[key]
+            if key == "city":
+                value = City.from_json(value)
+            setattr(self, key, value)
+
+
+@dataclass
 class Place:
     """The city or airport"""
 
+    id: str
+    name: str
+    type: str
+    iata_city_code: Optional[str]
+    iata_country_code: str
+    latitude: Optional[float]
+    longitude: Optional[float]
+    icao_code: Optional[str]
+    time_zone: Optional[str]
+    city_name: Optional[str]
+    city: Optional[City]
+    airports: Optional[Sequence[CityAirport]]
+
     allowed_types = ["airport", "city"]
 
     class InvalidType(Exception):
         """Invalid type of place"""
 
-    def __init__(self, json):
-        for key in json:
-            value = json[key]
-            if key == "type" and value not in Place.allowed_types:
-                raise Place.InvalidType(value)
-            elif key == "city" and value:
-                value = City(value)
-            elif key == "airports" and value:
-                value = [CityAirport(v) for v in value]
-            setattr(self, key, value)
+    @classmethod
+    def from_json(cls, json: dict):
+        """Construct a class instance from a JSON response."""
 
+        type = json["type"]
+        if type not in Place.allowed_types:
 
-class CityAirport:
-    """The airport associated to a city."""
+            raise Place.InvalidType(type)
 
-    def __init__(self, json):
-        for key in json:
-            value = json[key]
-            if key == "city":
-                value = City(value)
-            setattr(self, key, value)
+        return cls(
+            id=json["id"],
+            name=json["name"],
+            type=json["type"],
+            iata_city_code=json.get("iata_city_code"),
+            iata_country_code=json["iata_country_code"],
+            latitude=json.get("latitude"),
+            longitude=json.get("longitude"),
+            icao_code=json.get("icao_code"),
+            time_zone=json.get("time_zone"),
+            city_name=json.get("city_name"),
+            city=get_and_transform(json, "city", City.from_json),
+            airports=get_and_transform(
+                json,
+                "airports",
+                lambda value: [CityAirport(airport) for airport in value],
+            ),
+        )

duffel_api/utils.py

@@ -49,6 +49,20 @@ def maybe_parse_date_entries(key: str, value: Any) -> Union[str, datetime, date]
     return value
 
 
+def identity(value: Any) -> Any:
+    """Given a value, return the exact same value"""
+    return value
+
+
+def get_and_transform(dict: dict, key: str, fn=identity):
+    """Get a value from a dictionary and transform it or return
+    None if the key isn't present"""
+    try:
+        return fn(dict[key])
+    except KeyError:
+        return None
+
+
 def version() -> str:
     """Return the version specified in the package (setup.py) during runtime"""
     import pkg_resources