Source code for spectral_indices.idb.models

from __future__ import annotations

import json
from importlib.resources import files
from typing import List, Type, TypeVar

from peewee import CharField, ForeignKeyField, Model, SqliteDatabase, TextField

INDEX_DB = SqliteDatabase(files("spectral_indices.data").joinpath("index.sqlite"))


M = TypeVar("M", bound="BaseModel")




class BaseModel(Model):
    """Base Model for IDB database."""

    class Meta:
        database = INDEX_DB



    @classmethod
    def get_or_raise(cls: Type[M], **kwargs) -> M:
        """Return object with kwargs query (from table attributes) or raise error if no object fit the query.

        Args:
            cls (``Type[M]``): A BaseModel.

        Raises:
            ValueError: No object of table {cls.__name__} with {kwargs=}

        Returns:
            ``M``:
                - Object of the table.
        """
        obj = cls.get_or_none(**kwargs)
        if not obj:
            raise ValueError(f"No object of table {cls.__name__} with {kwargs=}")

        return obj






class Sensor(BaseModel):
    """Sensor table."""

    name = CharField(unique=True)





class Index(BaseModel):
    """Index table"""

    short_name = CharField()
    long_name = CharField(unique=True)
    general_formula = CharField(max_length=100)
    origin = CharField(max_length=50)





class Formula(BaseModel):
    """Formula table"""

    sensor = ForeignKeyField(Sensor, on_delete="CASCADE")
    index = ForeignKeyField(Index, on_delete="CASCADE")
    formula = CharField(max_length=100, unique=False)
    bands = TextField(default="[]")



    def set_bands(self, bands: List[str]):
        """Set bands to Formula data. Store list as serialized object.

        Args:
            bands (``List[str]``): List of bands IDs.
        """
        self.bands = json.dumps(bands)




    def get_bands(self) -> List[str]:
        """Unserialized bands from Formula data.

        Returns:
            ``List[str]``:
                - List of band IDs.
        """
        return json.loads(self.bands)