Add type annotations for models

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
@@ -21,6 +22,7 @@
 from .order_change_request import OrderChangeRequest
 from .payment import Payment
 from .payment_intent import PaymentIntent
+from .refund import Refund
 from .seat_map import SeatMap
 from .webhook import Webhook
 
@@ -45,6 +47,7 @@
     "OrderChangeRequest",
     "Payment",
     "PaymentIntent",
+    "Refund",
     "SeatMap",
     "Webhook",
 ]

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/airline.py

@@ -1,9 +1,21 @@
+from dataclasses import dataclass
+
+
+@dataclass
 class Airline:
     """Airlines are used to identify the air travel companies selling and operating
     flights
-
     """
 
-    def __init__(self, json):
-        for key in json:
-            setattr(self, key, json[key])
+    id: str
+    name: str
+    iata_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"],
+        )

duffel_api/models/airport.py

@@ -1,21 +1,36 @@
-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
+from ..utils import get_and_transform
 
-    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: Optional[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=get_and_transform(json, "city", City.from_json),
+        )

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/loyalty_programme_account.py

@@ -1,6 +1,19 @@
+
+
+from dataclasses import dataclass
+
+
+@dataclass
 class LoyaltyProgrammeAccount:
     """A passenger's loyalty programme account"""
 
-    def __init__(self, json):
-        for key in json:
-            setattr(self, key, json[key])
+    airline_iata_code: str
+    account_number: str
+
+    @classmethod
+    def from_json(cls, json: dict):
+        """Construct a class instance from a JSON response."""
+        return cls(
+            airline_iata_code=json["airline_iata_code"],
+            account_number=json["account_number"],
+        )

duffel_api/models/order_cancellation.py

@@ -1,6 +1,8 @@
-from ..utils import maybe_parse_date_entries
+from dataclasses import dataclass
+from datetime import datetime
 
 
+@dataclass
 class OrderCancellation:
     """To cancel an order, you'll need to create an order cancellation,
     check the refund_amount returned, and, if you're happy to go ahead and
@@ -11,6 +13,16 @@ class OrderCancellation:
     You'll then need to refund your customer (e.g. back to their credit/debit card).
     """
 
+    id: str
+    order_id: str
+    live_mode: bool
+    expires_at: datetime
+    refund_amount: str
+    refund_currency: str
+    refund_to: str
+    confirmed_at: datetime
+    created_at: datetime
+
     allowed_refund_types = [
         "arc_bsp_cash",
         "balance",
@@ -22,12 +34,32 @@ class OrderCancellation:
     class InvalidRefundType(Exception):
         """Invalid refund type provided"""
 
-    def __init__(self, json):
-        for key in json:
-            value = maybe_parse_date_entries(key, json[key])
-            if (
-                key == "refund_to"
-                and value not in OrderCancellation.allowed_refund_types
-            ):
-                raise OrderCancellation.InvalidRefundType(value)
-            setattr(self, key, value)
+    def __post_init__(self):
+        if self.refund_to not in OrderCancellation.allowed_refund_types:
+            raise OrderCancellation.InvalidRefundType(self.refund_to)
+
+        return self
+
+    @classmethod
+    def from_json(cls, json: dict):
+        """Construct a class instance from a JSON response."""
+        return cls(
+            id=json["id"],
+            order_id=json["order_id"],
+            live_mode=json["live_mode"],
+            expires_at=parse_datetime(json["confirmed_at"]),
+            refund_amount=json["refund_amount"],
+            refund_currency=json["refund_currency"],
+            refund_to=json["refund_to"],
+            confirmed_at=parse_datetime(json["confirmed_at"]),
+            created_at=datetime.strptime(json["created_at"], "%Y-%m-%dT%H:%M:%S.%fZ"),
+        )
+
+
+def parse_datetime(value: str) -> datetime:
+    # There are inconsistent formats used for this field depending on the
+    # endpoint
+    if len(value) == 20:
+        return datetime.strptime(value, "%Y-%m-%dT%H:%M:%SZ")
+    else:
+        return datetime.strptime(value, "%Y-%m-%dT%H:%M:%S.%fZ")

duffel_api/models/payment_intent.py

@@ -1,6 +1,12 @@
-from ..utils import maybe_parse_date_entries
+from dataclasses import dataclass
+from datetime import datetime
+from typing import Optional, Sequence
 
+from ..models import Refund
+from ..utils import get_and_transform
 
+
+@dataclass
 class PaymentIntent:
     """To begin the process of collecting a card payment from your customer, you
     need to create a Payment Intent.
@@ -11,7 +17,46 @@ class PaymentIntent:
     If the Payment Intent is created in test mode you should use a test card.
     """
 
-    def __init__(self, json):
-        for key in json:
-            value = maybe_parse_date_entries(key, json[key])
-            setattr(self, key, value)
+    id: str
+    live_mode: bool
+    amount: str
+    currency: str
+    net_amount: Optional[str]
+    net_currency: Optional[str]
+    fees_amount: Optional[str]
+    fees_currency: Optional[str]
+    client_token: str
+    card_network: Optional[str]
+    card_last_four_digits: Optional[str]
+    card_country_code: Optional[str]
+    status: str
+    refunds: Sequence[Refund]
+    created_at: datetime
+    updated_at: datetime
+
+    @classmethod
+    def from_json(cls, json: dict):
+        """Construct a class instance from a JSON response."""
+        return cls(
+            id=json["id"],
+            live_mode=json["live_mode"],
+            amount=json["amount"],
+            currency=json["currency"],
+            net_amount=json.get("net_amount"),
+            net_currency=json.get("net_currency"),
+            fees_amount=json.get("fees_amount"),
+            fees_currency=json.get("fees_currency"),
+            client_token=json["client_token"],
+            card_network=json.get("card_network"),
+            card_last_four_digits=json.get("card_last_four_digits"),
+            card_country_code=json.get("card_country_code"),
+            status=json["status"],
+            refunds=get_and_transform(
+                json,
+                "refunds",
+                lambda value: [Refund.from_json(refund) for refund in value],
+                []
+            ),
+            created_at=datetime.strptime(json["created_at"], "%Y-%m-%dT%H:%M:%S.%fZ"),
+            updated_at=datetime.strptime(json["created_at"], "%Y-%m-%dT%H:%M:%S.%fZ"),
+        )

duffel_api/models/place.py

@@ -1,32 +1,56 @@
-from ..models import City
+from dataclasses import dataclass
+from typing import Optional, Sequence
 
+from ..models import Airport, City
+from ..utils import get_and_transform
 
+
+@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[Airport]]
+
     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)
-
-
-class CityAirport:
-    """The airport associated to a city."""
-
-    def __init__(self, json):
-        for key in json:
-            value = json[key]
-            if key == "city":
-                value = City(value)
-            setattr(self, key, value)
+    def __post_init__(self):
+        if self.type not in Place.allowed_types:
+            raise Place.InvalidType(self.type)
+
+        return self
+
+    @classmethod
+    def from_json(cls, json: dict):
+        """Construct a class instance from a JSON response."""
+        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: [Airport.from_json(airport) for airport in value],
+            ),
+        )