Skip to content
This repository has been archived by the owner on Sep 12, 2024. It is now read-only.

Commit

Permalink
Add type annotations for models
Browse files Browse the repository at this point in the history
  • Loading branch information
Jesse Claven committed Jan 17, 2022
1 parent 18d418c commit 38281bc
Show file tree
Hide file tree
Showing 12 changed files with 340 additions and 72 deletions.
5 changes: 4 additions & 1 deletion duffel_api/models/__init__.py
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -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

Expand All @@ -45,6 +47,7 @@
"OrderChangeRequest",
"Payment",
"PaymentIntent",
"Refund",
"SeatMap",
"Webhook",
]
19 changes: 16 additions & 3 deletions duffel_api/models/aircraft.py
Original file line number Diff line number Diff line change
@@ -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"],
)
20 changes: 16 additions & 4 deletions duffel_api/models/airline.py
Original file line number Diff line number Diff line change
@@ -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"],
)
49 changes: 32 additions & 17 deletions duffel_api/models/airport.py
Original file line number Diff line number Diff line change
@@ -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),
)
24 changes: 24 additions & 0 deletions duffel_api/models/city.py
Original file line number Diff line number Diff line change
@@ -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"],
)
19 changes: 16 additions & 3 deletions duffel_api/models/loyalty_programme_account.py
Original file line number Diff line number Diff line change
@@ -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"],
)
52 changes: 42 additions & 10 deletions duffel_api/models/order_cancellation.py
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -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",
Expand All @@ -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")
54 changes: 49 additions & 5 deletions duffel_api/models/payment_intent.py
Original file line number Diff line number Diff line change
@@ -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.
Expand All @@ -11,7 +17,45 @@ 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"),
)
68 changes: 46 additions & 22 deletions duffel_api/models/place.py
Original file line number Diff line number Diff line change
@@ -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],
),
)
Loading

0 comments on commit 38281bc

Please sign in to comment.