add flight.py and test

Author: matthiaszimmermann

examples/flight.py

@@ -0,0 +1,91 @@
+"""Example of using Pydantic with pytz for flight time handling."""
+
+from datetime import datetime
+from typing import Optional
+
+from pydantic import BaseModel, Field, field_validator
+import pytz
+
+
+class Flight(BaseModel):
+    """Represents a flight with departure and arrival times in different timezones."""
+
+    flight_number: str = Field(..., pattern=r"^[A-Z]{2}\d{3,4}$")
+    departure_airport: str = Field(..., min_length=3, max_length=3)
+    arrival_airport: str = Field(..., min_length=3, max_length=3)
+    departure_time: datetime
+    arrival_time: datetime
+    departure_timezone: str = Field(..., pattern=r"^[A-Za-z]+/[A-Za-z_]+$")
+    arrival_timezone: str = Field(..., pattern=r"^[A-Za-z]+/[A-Za-z_]+$")
+    duration_minutes: Optional[int] = None
+
+    @field_validator("departure_timezone", "arrival_timezone")
+    @classmethod
+    def validate_timezone(cls, v: str) -> str:
+        """Validate that the timezone exists in pytz."""
+        try:
+            pytz.timezone(v)
+        except pytz.exceptions.UnknownTimeZoneError:
+            raise ValueError(f"Unknown timezone: {v}")
+        return v
+
+    @field_validator("arrival_time")
+    @classmethod
+    def validate_arrival_time(cls, v: datetime, info) -> datetime:
+        """Validate that arrival time is after departure time."""
+        departure_time = info.data.get("departure_time")
+        if departure_time and v <= departure_time:
+            raise ValueError("Arrival time must be after departure time")
+        return v
+
+    def calculate_duration(self) -> int:
+        """Calculate flight duration in minutes."""
+        if self.duration_minutes is not None:
+            return self.duration_minutes
+
+        # Convert times to UTC for comparison
+        departure_tz = pytz.timezone(self.departure_timezone)
+        arrival_tz = pytz.timezone(self.arrival_timezone)
+
+        departure_utc = departure_tz.localize(self.departure_time).astimezone(pytz.UTC)
+        arrival_utc = arrival_tz.localize(self.arrival_time).astimezone(pytz.UTC)
+
+        duration = arrival_utc - departure_utc
+        return int(duration.total_seconds() / 60)
+
+    def to_local_times(self) -> dict[str, datetime]:
+        """Convert times to their respective timezones."""
+        departure_tz = pytz.timezone(self.departure_timezone)
+        arrival_tz = pytz.timezone(self.arrival_timezone)
+
+        return {
+            "departure_local": departure_tz.localize(self.departure_time),
+            "arrival_local": arrival_tz.localize(self.arrival_time),
+        }
+
+
+def main() -> None:
+    """Example usage of the Flight class."""
+    # Example flight from New York to London
+    flight = Flight(
+        flight_number="BA178",
+        departure_airport="JFK",
+        arrival_airport="LHR",
+        departure_time=datetime(2024, 3, 25, 20, 0),  # 8 PM
+        arrival_time=datetime(2024, 3, 26, 8, 0),  # 8 AM next day
+        departure_timezone="America/New_York",
+        arrival_timezone="Europe/London",
+    )
+
+    # Calculate duration
+    duration = flight.calculate_duration()
+    print(f"Flight duration: {duration} minutes")
+
+    # Get local times
+    local_times = flight.to_local_times()
+    print(f"Departure (local): {local_times['departure_local']}")
+    print(f"Arrival (local): {local_times['arrival_local']}")
+
+
+if __name__ == "__main__":
+    main()

tests/test_flight.py

@@ -0,0 +1,116 @@
+"""Tests for the Flight class."""
+
+from datetime import datetime
+
+import pytest
+from examples.flight import Flight
+
+
+def test_valid_flight() -> None:
+    """Test creating a valid flight."""
+    flight = Flight(
+        flight_number="BA178",
+        departure_airport="JFK",
+        arrival_airport="LHR",
+        departure_time=datetime(2024, 3, 25, 20, 0),
+        arrival_time=datetime(2024, 3, 26, 8, 0),
+        departure_timezone="America/New_York",
+        arrival_timezone="Europe/London",
+    )
+
+    assert flight.flight_number == "BA178"
+    assert flight.departure_airport == "JFK"
+    assert flight.arrival_airport == "LHR"
+    assert flight.duration_minutes is None
+
+
+def test_invalid_flight_number() -> None:
+    """Test invalid flight number format."""
+    with pytest.raises(ValueError):
+        Flight(
+            flight_number="123",  # Invalid format
+            departure_airport="JFK",
+            arrival_airport="LHR",
+            departure_time=datetime(2024, 3, 25, 20, 0),
+            arrival_time=datetime(2024, 3, 26, 8, 0),
+            departure_timezone="America/New_York",
+            arrival_timezone="Europe/London",
+        )
+
+
+def test_invalid_airport_code() -> None:
+    """Test invalid airport code length."""
+    with pytest.raises(ValueError):
+        Flight(
+            flight_number="BA178",
+            departure_airport="JFKX",  # Too long
+            arrival_airport="LHR",
+            departure_time=datetime(2024, 3, 25, 20, 0),
+            arrival_time=datetime(2024, 3, 26, 8, 0),
+            departure_timezone="America/New_York",
+            arrival_timezone="Europe/London",
+        )
+
+
+def test_invalid_timezone() -> None:
+    """Test invalid timezone."""
+    with pytest.raises(ValueError):
+        Flight(
+            flight_number="BA178",
+            departure_airport="JFK",
+            arrival_airport="LHR",
+            departure_time=datetime(2024, 3, 25, 20, 0),
+            arrival_time=datetime(2024, 3, 26, 8, 0),
+            departure_timezone="Invalid/Timezone",
+            arrival_timezone="Europe/London",
+        )
+
+
+def test_arrival_before_departure() -> None:
+    """Test that arrival time must be after departure time."""
+    with pytest.raises(ValueError):
+        Flight(
+            flight_number="BA178",
+            departure_airport="JFK",
+            arrival_airport="LHR",
+            departure_time=datetime(2024, 3, 25, 20, 0),
+            arrival_time=datetime(2024, 3, 25, 19, 0),  # Before departure
+            departure_timezone="America/New_York",
+            arrival_timezone="Europe/London",
+        )
+
+
+def test_duration_calculation() -> None:
+    """Test flight duration calculation."""
+    flight = Flight(
+        flight_number="BA178",
+        departure_airport="JFK",
+        arrival_airport="LHR",
+        departure_time=datetime(2024, 3, 25, 20, 0),
+        arrival_time=datetime(2024, 3, 26, 8, 0),
+        departure_timezone="America/New_York",
+        arrival_timezone="Europe/London",
+    )
+
+    duration = flight.calculate_duration()
+    assert duration > 0  # Should be positive
+    assert duration < 1440  # Should be less than 24 hours
+
+
+def test_local_times() -> None:
+    """Test local time conversion."""
+    flight = Flight(
+        flight_number="BA178",
+        departure_airport="JFK",
+        arrival_airport="LHR",
+        departure_time=datetime(2024, 3, 25, 20, 0),
+        arrival_time=datetime(2024, 3, 26, 8, 0),
+        departure_timezone="America/New_York",
+        arrival_timezone="Europe/London",
+    )
+
+    local_times = flight.to_local_times()
+    assert "departure_local" in local_times
+    assert "arrival_local" in local_times
+    assert local_times["departure_local"].tzinfo is not None
+    assert local_times["arrival_local"].tzinfo is not None