Source code for kappa_sdk.shutin_list

from __future__ import annotations
from .shutin import ShutIn
from typing import List, Optional, Iterator
from .data import Data
from .data_kind_enum import DataKindEnum




class ShutInList(list):  # type: ignore[type-arg]

    """
    Class to store a List of Shut-in
    """

    def __init__(self, shutins: Optional[List[ShutIn]] = None) -> None:
        if shutins is None:
            shutins = []
        super().__init__(shutins)

    def __iter__(self) -> Iterator[ShutIn]:
        return iter(list.__iter__(self))

    @staticmethod
    def from_data(data: Data) -> ShutInList:
        """ Initialize a :class:`ShutInList` object from a :class:`Data` object

        Parameters
        ----------
        data:
            a :class:`Data` object

        Returns
        -------
        :class:`ShutInList`
            A list of :class:`ShutIn` objects
        """
        if data.data_type == 'SI':
            shutin_attributes = data._Data__cluster_apis.data_api.get_shutin_attributes(data.vector_id)  # type: ignore[attr-defined]
            pressure_id = data.get_shutin_data_dto().input.xyVectorId if data.kind == DataKindEnum.shutin else None
            shutins: List[ShutIn] = list()
            for shutin in shutin_attributes.attributes:
                shutins.append(ShutIn(shutin.start, shutin.end, shutin.attributes.reference, pressure_id, shutin.attributes.category.value, shutin.attributes.type.value, data._Data__cluster_apis.kw_api))  # type: ignore[attr-defined]
            return ShutInList(shutins)
        else:
            raise ValueError(f"ShutInList object can only be initialize with {DataKindEnum.shutin} data")