feat: parameterise fields in livefeed (#30)

cathaypacific8747 · Jun 18, 2024 · 8965fe2 · 8965fe2
1 parent 2f6223e
commit 8965fe2
Show file tree

Hide file tree

Showing 7 changed files with 103 additions and 38 deletions.
diff --git a/docs/usage/scripts/12_livefeed.py b/docs/usage/scripts/12_livefeed.py
@@ -42,6 +42,7 @@ async def my_feed() -> None:
  "destination": "AKL",
  "typecode": "B744",
  "eta": 0,
+ "squawk": 0,
  }
  # ... 15109 more items
  ],
@@ -50,18 +51,18 @@ async def my_feed() -> None:
 #%%
 """
 # --8<-- [start:df]
- timestamp flightid latitude longitude heading altitude ground_speed on_ground callsign source registration origin destination typecode eta vertical_speed
-0 1711911905 882151247 -12.432657 -172.148254 203 34000 515 False QFA7552 0 N409MC HNL AKL B744 0 0
-1 1711911902 882203620 -16.504490 -178.940308 249 36000 460 False VOZ76 0 VH-YIL APW BNE B738 0 0
-2 1711911904 882212062 -13.240505 -176.195602 292 0 0 True RLY0100 0 F-OCQZ WLS DHC6 0 0
-3 1711911897 882145424 10.347591 -167.007263 56 37000 516 False QFA15 5 VH-EBQ BNE LAX A332 0 0
-4 1711911905 882199081 18.591330 -165.391083 247 32000 416 False UAL132 3 N77296 HNL MAJ B738 0 0
-... ... ... ... ... ... ... ... ... ... ... ... ... ... ... ... ...
-15095 1711911897 882140495 53.746582 174.467392 258 38000 435 False GTI8650 5 N710GT LAX HKG B77L 0 0
-15096 1711911887 882155955 54.291172 175.606842 57 31000 515 False UPS81 5 N628UP PVG ANC B748 0 0
-15097 1711911673 882165589 57.548691 179.460800 246 36000 438 False KAL258 4 HL8285 ANC ICN B77L 0 0
-15098 1711911901 882187089 56.527115 176.298798 242 38000 442 False CKS936 4 N701CK ANC HFE B744 0 0
-15099 1711911905 882137160 59.153641 179.730972 229 38000 471 False KAL8286 5 HL8043 YYZ ICN B77L 0 0
+ timestamp flightid latitude longitude heading altitude ground_speed on_ground callsign source registration origin destination typecode eta vertical_speed squawk
+0 1711911905 882151247 -12.432657 -172.148254 203 34000 515 False QFA7552 0 N409MC HNL AKL B744 0 0 0
+1 1711911902 882203620 -16.504490 -178.940308 249 36000 460 False VOZ76 0 VH-YIL APW BNE B738 0 0 0
+2 1711911904 882212062 -13.240505 -176.195602 292 0 0 True RLY0100 0 F-OCQZ WLS DHC6 0 0 0
+3 1711911897 882145424 10.347591 -167.007263 56 37000 516 False QFA15 5 VH-EBQ BNE LAX A332 0 0 0
+4 1711911905 882199081 18.591330 -165.391083 247 32000 416 False UAL132 3 N77296 HNL MAJ B738 0 0 0
+... ... ... ... ... ... ... ... ... ... ... ... ... ... ... ... ... ...
+15095 1711911897 882140495 53.746582 174.467392 258 38000 435 False GTI8650 5 N710GT LAX HKG B77L 0 0 0
+15096 1711911887 882155955 54.291172 175.606842 57 31000 515 False UPS81 5 N628UP PVG ANC B748 0 0 0
+15097 1711911673 882165589 57.548691 179.460800 246 36000 438 False KAL258 4 HL8285 ANC ICN B77L 0 0 0
+15098 1711911901 882187089 56.527115 176.298798 242 38000 442 False CKS936 4 N701CK ANC HFE B744 0 0 0
+15099 1711911905 882137160 59.153641 179.730972 229 38000 471 False KAL8286 5 HL8043 YYZ ICN B77L 0 0 0
 
 [15100 rows x 16 columns]
 # --8<-- [end:df]
@@ -83,18 +84,18 @@ async def my_feed() -> None:
 # %%
 """
 # --8<-- [start:df2]
- timestamp flightid latitude longitude heading altitude ground_speed on_ground callsign source registration origin destination typecode eta vertical_speed
-0 1711697713 881642343 -13.628512 -169.556046 44 37000 499 False DAL42 0 N519DN SYD LAX A359 0 0
-1 1711697709 881686001 -14.096649 -172.093018 182 8850 298 False ANZ999 0 ZK-NNE APW AKL A21N 0 0
-2 1711697713 881673319 -46.344315 -173.242355 122 31975 568 False LAN800 4 CC-BGC AKL SCL B789 0 0
-3 1711697713 881679659 -34.398090 -178.388153 64 37975 530 False ANZ902 4 ZK-NZE AKL PPT B789 0 0
-4 1711697717 881655685 -12.116470 -169.877808 199 36000 469 False HAL465 5 N395HA HNL PPG A332 0 0
-... ... ... ... ... ... ... ... ... ... ... ... ... ... ... ... ...
-11422 1711697721 881662489 56.048801 174.987000 240 36000 468 False CPA3187 3 B-LJG ANC HKG B748 0 0
-11423 1711697713 881602916 56.870163 175.223923 226 35000 476 False ANA111 3 JA796A ORD HND B77W 0 0
-11424 1711697713 881669755 56.652039 176.742554 240 34000 459 False FDX5236 3 N573FE ANC ICN MD11 0 0
-11425 1711697721 881652749 57.509033 179.559845 242 38000 467 False GTI8065 3 N452PA PDX ICN B744 0 0
-11426 1711697717 881648986 54.571342 175.497025 55 31000 501 False EVA662 5 B-16781 TPE ANC B77L 0 0
+ timestamp flightid latitude longitude heading altitude ground_speed on_ground callsign source registration origin destination typecode eta vertical_speed squawk
+0 1711697713 881642343 -13.628512 -169.556046 44 37000 499 False DAL42 0 N519DN SYD LAX A359 0 0 0
+1 1711697709 881686001 -14.096649 -172.093018 182 8850 298 False ANZ999 0 ZK-NNE APW AKL A21N 0 0 0
+2 1711697713 881673319 -46.344315 -173.242355 122 31975 568 False LAN800 4 CC-BGC AKL SCL B789 0 0 0
+3 1711697713 881679659 -34.398090 -178.388153 64 37975 530 False ANZ902 4 ZK-NZE AKL PPT B789 0 0 0
+4 1711697717 881655685 -12.116470 -169.877808 199 36000 469 False HAL465 5 N395HA HNL PPG A332 0 0 0
+... ... ... ... ... ... ... ... ... ... ... ... ... ... ... ... ... ...
+11422 1711697721 881662489 56.048801 174.987000 240 36000 468 False CPA3187 3 B-LJG ANC HKG B748 0 0 0
+11423 1711697713 881602916 56.870163 175.223923 226 35000 476 False ANA111 3 JA796A ORD HND B77W 0 0 0
+11424 1711697713 881669755 56.652039 176.742554 240 34000 459 False FDX5236 3 N573FE ANC ICN MD11 0 0 0
+11425 1711697721 881652749 57.509033 179.559845 242 38000 467 False GTI8065 3 N452PA PDX ICN B744 0 0 0
+11426 1711697717 881648986 54.571342 175.497025 55 31000 501 False EVA662 5 B-16781 TPE ANC B77L 0 0 0
 
 [11427 rows x 16 columns]
 # --8<-- [end:df2]

diff --git a/src/fr24/core.py b/src/fr24/core.py
@@ -41,6 +41,7 @@
 from .types.fr24 import (
  FlightData,
  FlightList,
+ LivefeedField,
  Playback,
  TokenSubscriptionKey,
  UsernamePassword,
@@ -469,19 +470,26 @@ async def _fetch(
  self,
  ctx: LiveFeedContext,
  ) -> list[LiveFeedRecord]:
+ kw = {
+ k: v
+ for k, v in ctx.items()
+ if k in ("limit", "fields") and v is not None
+ }
  if (ts := ctx["timestamp"]) is not None:
- kw = {
- k: v
- for k, v in ctx.items()
- if k in ("duration", "hfreq") and v is not None
- }
+ kw.update(
+ {
+ k: v
+ for k, v in ctx.items()
+ if k in ("duration", "hfreq") and v is not None
+ }
+ )
  return await livefeed_playback_world_data(
  self.http.client,
  ts,
  **kw, # type: ignore[arg-type]
  auth=self.http.auth,
  )
- resp = await livefeed_world_data(self.http.client, self.http.auth)
+ resp = await livefeed_world_data(self.http.client, self.http.auth, **kw) # type: ignore[arg-type]
  ctx["timestamp"] = int(time.time())
  # TODO: use server time instead, but it doesn't really matter because
  # livefeed messages have timestamps attached to them anyway
@@ -556,6 +564,13 @@ async def fetch(
  *,
  duration: int | None = None,
  hfreq: int | None = None,
+ limit: int = 1500,
+ fields: list[LivefeedField] = [
+ "flight",
+ "reg",
+ "route",
+ "type",
+ ],
  ) -> LiveFeedAPIResp:
  """
  Fetch live feed data.
@@ -569,8 +584,13 @@ async def fetch(
  be set for historical data.
  :param hfreq: High frequency mode (default: `0`). Should only be set
  for historical data.
+ :param limit: Max number of flights (default 1500 for unauthenticated
+ users, 2000 for authenticated users)
+ :param fields: fields to include - for unauthenticated users, max 4
+ fields. When authenticated, `squawk`, `vspeed`, `airspace`,
+ `logo_id` and `age` can be included
  """
- ctx = self._construct_ctx(timestamp, duration, hfreq)
+ ctx = self._construct_ctx(timestamp, duration, hfreq, limit, fields)
  return LiveFeedAPIResp(ctx, await self._api._fetch(ctx))
 
  def load(
@@ -584,14 +604,16 @@ def load(
 
  :param timestamp: Unix timestamp (seconds) of the saved feed snapshot.
  """
- ctx = self._construct_ctx(timestamp, None, None)
+ ctx = self._construct_ctx(timestamp, None, None, None, None)
  return LiveFeedArrow.from_cache(ctx)
 
  def _construct_ctx(
  self,
  timestamp: int | str | datetime | pd.Timestamp | None,
  duration: int | None,
  hfreq: int | None,
+ limit: int | None,
+ fields: list[LivefeedField] | None,
  ) -> LiveFeedContext:
  ts = to_unix_timestamp(timestamp)
  if ts is None and (hfreq is not None or duration is not None):
@@ -603,5 +625,7 @@ def _construct_ctx(
  "source": "live" if ts is None else "playback",
  "duration": duration,
  "hfreq": hfreq,
+ "limit": limit,
+ "fields": fields,
  "base_dir": self._base_dir,
  }
diff --git a/src/fr24/livefeed.py b/src/fr24/livefeed.py
@@ -22,7 +22,7 @@
  VisibilitySettings,
 )
 from .types.cache import LiveFeedRecord
-from .types.fr24 import Authentication
+from .types.fr24 import Authentication, LivefeedField
 
 # N, S, W, E
 world_zones = [
@@ -39,7 +39,7 @@ def livefeed_message_create(
  stats: bool = False,
  limit: int = 1500,
  maxage: int = 14400,
- fields: list[str] = [
+ fields: list[LivefeedField] = [
  "flight",
  "reg",
  "route",
@@ -193,20 +193,32 @@ def livefeed_flightdata_dict(
  "destination": lfr.extra_info.route.to,
  "typecode": lfr.extra_info.type,
  "eta": lfr.extra_info.schedule.eta,
+ "squawk": lfr.extra_info.squawk,
  }
 
 
 # TODO: add parameter for custom bounds, e.g. from .bounds.lng_bounds_per_30_min
 async def livefeed_world_data(
- client: httpx.AsyncClient, auth: None | Authentication = None
+ client: httpx.AsyncClient,
+ auth: None | Authentication = None,
+ limit: int = 1500,
+ fields: list[LivefeedField] = [
+ "flight",
+ "reg",
+ "route",
+ "type",
+ ],
 ) -> list[LiveFeedRecord]:
  """Retrieve live feed data for the entire world, in chunks."""
  results = await asyncio.gather(
  *[
  livefeed_post(
  client,
  livefeed_request_create(
- livefeed_message_create(*bounds), auth=auth
+ livefeed_message_create(
+ *bounds, limit=limit, fields=fields
+ ),
+ auth=auth,
  ),
  )
  for bounds in world_zones
@@ -225,6 +237,13 @@ async def livefeed_playback_world_data(
  duration: int = 7,
  hfreq: int = 0,
  auth: None | Authentication = None,
+ limit: int = 1500,
+ fields: list[LivefeedField] = [
+ "flight",
+ "reg",
+ "route",
+ "type",
+ ],
 ) -> list[LiveFeedRecord]:
  """
  Retrieve live feed playback data for the entire world, in chunks.
@@ -237,7 +256,9 @@ async def livefeed_playback_world_data(
  client,
  livefeed_playback_request_create(
  livefeed_playback_message_create(
- livefeed_message_create(*bounds),
+ livefeed_message_create(
+ *bounds, limit=limit, fields=fields
+ ),
  timestamp,
  timestamp + duration,
  hfreq,

diff --git a/src/fr24/types/cache.py b/src/fr24/types/cache.py
@@ -137,6 +137,7 @@ class PlaybackTrackEMSRecord(TypedDict):
  pa.field("typecode", pa.string()),
  pa.field("eta", pa.uint32()),
  pa.field("vertical_speed", pa.int16()), # 64 * 9-bit + 1-bit sign
+ pa.field("squawk", pa.uint16()),
  ]
 )
 
@@ -158,3 +159,4 @@ class LiveFeedRecord(TypedDict):
  destination: str
  typecode: str
  eta: int
+ squawk: int
diff --git a/src/fr24/types/core.py b/src/fr24/types/core.py
@@ -5,6 +5,8 @@
 
 from typing_extensions import NotRequired
 
+from .fr24 import LivefeedField
+
 
 # internally tagged
 class FlightListContext(TypedDict):
@@ -23,4 +25,6 @@ class LiveFeedContext(TypedDict):
  source: Literal["live", "playback"]
  duration: NotRequired[int | None]
  hfreq: NotRequired[int | None]
+ limit: NotRequired[int | None]
+ fields: NotRequired[list[LivefeedField] | None]
  base_dir: NotRequired[Path]
diff --git a/src/fr24/types/fr24.py b/src/fr24/types/fr24.py
@@ -477,3 +477,16 @@ class AirportResult(TypedDict):
 class AirportList(TypedDict):
  result: AirportResult
  _api: APIResult
+
+
+LivefeedField = Literal[
+ "flight",
+ "reg",
+ "route",
+ "type",
+ "squawk",
+ "vspeed",
+ "airspace",
+ "logo_id",
+ "age",
+]
diff --git a/tests/conftest.py b/tests/conftest.py
@@ -9,4 +9,4 @@
 @pytest.fixture(autouse=True)
 def slow_down():
  yield
- time.sleep(2) # avoid overloading the server
+ time.sleep(2.5) # avoid overloading the server