core/homeassistant/components/seventeentrack/coordinator.py

"""Coordinator for 17Track."""

from dataclasses import dataclass
from typing import Any

from pyseventeentrack import Client as SeventeenTrackClient
from pyseventeentrack.errors import SeventeenTrackError
from pyseventeentrack.package import Package

from homeassistant.config_entries import ConfigEntry
from homeassistant.core import HomeAssistant
from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
from homeassistant.util import slugify

from .const import (
    CONF_SHOW_ARCHIVED,
    CONF_SHOW_DELIVERED,
    DEFAULT_SCAN_INTERVAL,
    DOMAIN,
    LOGGER,
)


@dataclass
class SeventeenTrackData:
    """Class for handling the data retrieval."""

    summary: dict[str, dict[str, Any]]
    live_packages: dict[str, Package]


class SeventeenTrackCoordinator(DataUpdateCoordinator[SeventeenTrackData]):
    """Class to manage fetching 17Track data."""

    config_entry: ConfigEntry

    def __init__(self, hass: HomeAssistant, client: SeventeenTrackClient) -> None:
        """Initialize."""
        super().__init__(
            hass,
            LOGGER,
            name=DOMAIN,
            update_interval=DEFAULT_SCAN_INTERVAL,
        )
        self.show_delivered = self.config_entry.options[CONF_SHOW_DELIVERED]
        self.account_id = client.profile.account_id

        self.show_archived = self.config_entry.options[CONF_SHOW_ARCHIVED]
        self.client = client

    async def _async_update_data(self) -> SeventeenTrackData:
        """Fetch data from 17Track API."""

        try:
            summary = await self.client.profile.summary(
                show_archived=self.show_archived
            )

            live_packages = set(
                await self.client.profile.packages(show_archived=self.show_archived)
            )

        except SeventeenTrackError as err:
            raise UpdateFailed(err) from err

        summary_dict = {}
        live_packages_dict = {}

        for status, quantity in summary.items():
            summary_dict[slugify(status)] = {
                "quantity": quantity,
                "packages": [],
                "status_name": status,
            }

        for package in live_packages:
            live_packages_dict[package.tracking_number] = package
            summary_value = summary_dict.get(slugify(package.status))
            if summary_value:
                summary_value["packages"].append(package)

        return SeventeenTrackData(
            summary=summary_dict, live_packages=live_packages_dict
        )