sensor.py

Go to the documentation of this file.

"""Support for non-delivered packages recorded in AfterShip."""
 
from __future__ import annotations
 
import logging
from typing import Any, Final
 
from pyaftership import AfterShip, AfterShipException
 
from homeassistant.components.sensor import SensorEntity
from homeassistant.core import HomeAssistant, ServiceCall
import homeassistant.helpers.config_validation as cv
from homeassistant.helpers.dispatcher import (
    async_dispatcher_connect,
    async_dispatcher_send,
)
from homeassistant.helpers.entity_platform import AddEntitiesCallback
from homeassistant.util import Throttle
 
from . import AfterShipConfigEntry
from .const import (
    ADD_TRACKING_SERVICE_SCHEMA,
    ATTR_TRACKINGS,
    ATTRIBUTION,
    BASE,
    CONF_SLUG,
    CONF_TITLE,
    CONF_TRACKING_NUMBER,
    DOMAIN,
    MIN_TIME_BETWEEN_UPDATES,
    REMOVE_TRACKING_SERVICE_SCHEMA,
    SERVICE_ADD_TRACKING,
    SERVICE_REMOVE_TRACKING,
    UPDATE_TOPIC,
)
 
_LOGGER: Final = logging.getLogger(__name__)
 
PLATFORM_SCHEMA: Final = cv.removed(DOMAIN, raise_if_present=False)
 
 
async def async_setup_entry(
    hass: HomeAssistant,
    config_entry: AfterShipConfigEntry,
    async_add_entities: AddEntitiesCallback,
) -> None:
    """Set up AfterShip sensor entities based on a config entry."""
    aftership = config_entry.runtime_data
 
    async_add_entities([AfterShipSensor(aftership, config_entry.title)], True)
 
    async def handle_add_tracking(call: ServiceCall) -> None:
        """Call when a user adds a new Aftership tracking from Home Assistant."""
        await aftership.trackings.add(
            tracking_number=call.data[CONF_TRACKING_NUMBER],
            title=call.data.get(CONF_TITLE),
            slug=call.data.get(CONF_SLUG),
        )
        async_dispatcher_send(hass, UPDATE_TOPIC)
 
    hass.services.async_register(
        DOMAIN,
        SERVICE_ADD_TRACKING,
        handle_add_tracking,
        schema=ADD_TRACKING_SERVICE_SCHEMA,
    )
 
    async def handle_remove_tracking(call: ServiceCall) -> None:
        """Call when a user removes an Aftership tracking from Home Assistant."""
        await aftership.trackings.remove(
            tracking_number=call.data[CONF_TRACKING_NUMBER],
            slug=call.data[CONF_SLUG],
        )
        async_dispatcher_send(hass, UPDATE_TOPIC)
 
    hass.services.async_register(
        DOMAIN,
        SERVICE_REMOVE_TRACKING,
        handle_remove_tracking,
        schema=REMOVE_TRACKING_SERVICE_SCHEMA,
    )
 
 
class AfterShipSensor(SensorEntity):
    """Representation of a AfterShip sensor."""
 
    _attr_attribution = ATTRIBUTION
    _attr_native_unit_of_measurement: str = "packages"
    _attr_translation_key = "packages"
 
    def __init__(self, aftership: AfterShip, name: str) -> None:
        """Initialize the sensor."""
        self._attributes_attributes: dict[str, Any] = {}
        self._state_state: int | None = None
        self.aftershipaftership = aftership
        self._attr_name_attr_name = name
 
    @property
    def native_value(self) -> int | None:
        """Return the state of the sensor."""
        return self._state_state
 
    @property
    def extra_state_attributes(self) -> dict[str, str]:
        """Return attributes for the sensor."""
        return self._attributes_attributes
 
    async def async_added_to_hass(self) -> None:
        """Register callbacks."""
        self.async_on_removeasync_on_remove(
            async_dispatcher_connect(self.hasshass, UPDATE_TOPIC, self._force_update_force_update)
        )
 
    async def _force_update(self) -> None:
        """Force update of data."""
        await self.async_updateasync_update(no_throttle=True)
        self.async_write_ha_stateasync_write_ha_state()
 
    @Throttle(MIN_TIME_BETWEEN_UPDATES)
    async def async_update(self, **kwargs: Any) -> None:
        """Get the latest data from the AfterShip API."""
        try:
            trackings = await self.aftershipaftership.trackings.list()
        except AfterShipException as err:
            _LOGGER.error("Errors when querying AfterShip - %s", err)
            return
 
        status_to_ignore = {"delivered"}
        status_counts: dict[str, int] = {}
        parsed_trackings = []
        not_delivered_count = 0
 
        for track in trackings["trackings"]:
            status = track["tag"].lower()
            name = (
                track["tracking_number"] if track["title"] is None else track["title"]
            )
            last_checkpoint = (
                f"Shipment {track['tag'].lower()}"
                if not track["checkpoints"]
                else track["checkpoints"][-1]
            )
            status_counts[status] = status_counts.get(status, 0) + 1
            parsed_trackings.append(
                {
                    "name": name,
                    "tracking_number": track["tracking_number"],
                    "slug": track["slug"],
                    "link": f"{BASE}{track['slug']}/{track['tracking_number']}",
                    "last_update": track["updated_at"],
                    "expected_delivery": track["expected_delivery"],
                    "status": track["tag"],
                    "last_checkpoint": last_checkpoint,
                }
            )
 
            if status not in status_to_ignore:
                not_delivered_count += 1
            else:
                _LOGGER.debug("Ignoring %s as it has status: %s", name, status)
 
        self._attributes_attributes = {
            **status_counts,
            ATTR_TRACKINGS: parsed_trackings,
        }
 
        self._state_state = not_delivered_count