vessel_summary_result.py

from typing import List

import pandas as pd

from vortexasdk.logger import get_logger
from vortexasdk.api import VesselSummary
from vortexasdk.api.search_result import Result
from vortexasdk.result_conversions import create_dataframe, create_list

logger = get_logger(__name__)


class VesselSummaryResult(Result):
    """Container class that holds the result obtained from calling the `Vessel-Summary` endpoint."""

    def to_list(self) -> List[VesselSummary]:
        """Represent vessels as a list."""
        # noinspection PyTypeChecker
        return create_list(super().to_list(), VesselSummary)

    def to_df(self, columns=None) -> pd.DataFrame:
        """
        Represent vessels as a `pd.DataFrame`.

        # Arguments
            columns: The vessel features we want in the dataframe. Enter `columns='all'` to include all features.
            Defaults to `columns = ['vessel_id', 'timestamp', 'lat', 'lon', 'speed', 'draught', 'declared_destination', 'declared_eta']`.

        # Returns
        `pd.DataFrame` of vessels.

        """
        return create_dataframe(
            columns=columns,
            default_columns=DEFAULT_COLUMNS,
            data=super().to_list(),
            logger_description="Vessel Summary",
        )


DEFAULT_COLUMNS = [
    "vessel_id",
    "timestamp",
    "lat",
    "lon",
    "speed",
    "draught",
    "declared_destination",
    "declared_eta",
]