data_classes.py

"""Pydantic data models and other dataclasses. This is the only file that uses Optional[]
typing syntax instead of | None syntax to work with pydantic"""
from __future__ import annotations

import pathlib
import secrets
import shutil
from abc import ABC, abstractmethod
from enum import Enum, auto
from typing import TYPE_CHECKING, Any, List, Optional, Union

from fastapi import Request
from gradio_client.documentation import document
from gradio_client.utils import traverse

from . import wasm_utils

if not wasm_utils.IS_WASM or TYPE_CHECKING:
    from pydantic import BaseModel, RootModel, ValidationError
else:
    # XXX: Currently Pyodide V2 is not available on Pyodide,
    # so we install V1 for the Wasm version.
    from typing import Generic, TypeVar

    from pydantic import BaseModel as BaseModelV1
    from pydantic import ValidationError, schema_of

    # Map V2 method calls to V1 implementations.
    # Ref: https://docs.pydantic.dev/latest/migration/#changes-to-pydanticbasemodel
    class BaseModelMeta(type(BaseModelV1)):
        def __new__(cls, name, bases, dct):
            # Override `dct` to dynamically create a `Config` class based on `model_config`.
            if "model_config" in dct:
                config_class = type("Config", (), {})
                for key, value in dct["model_config"].items():
                    setattr(config_class, key, value)
                dct["Config"] = config_class
                del dct["model_config"]

            model_class = super().__new__(cls, name, bases, dct)
            return model_class

    class BaseModel(BaseModelV1, metaclass=BaseModelMeta):
        pass

    BaseModel.model_dump = BaseModel.dict  # type: ignore
    BaseModel.model_json_schema = BaseModel.schema  # type: ignore

    # RootModel is not available in V1, so we create a dummy class.
    PydanticUndefined = object()
    RootModelRootType = TypeVar("RootModelRootType")

    class RootModel(BaseModel, Generic[RootModelRootType]):
        root: RootModelRootType

        def __init__(self, root: RootModelRootType = PydanticUndefined, **data):
            if data:
                if root is not PydanticUndefined:
                    raise ValueError(
                        '"RootModel.__init__" accepts either a single positional argument or arbitrary keyword arguments'
                    )
                root = data  # type: ignore
            # XXX: No runtime validation is executed.
            super().__init__(root=root)  # type: ignore

        def dict(self, **kwargs):
            return super().dict(**kwargs)["root"]

        @classmethod
        def schema(cls, **_kwargs):
            # XXX: kwargs are ignored.
            return schema_of(cls.__fields__["root"].type_)  # type: ignore

    RootModel.model_dump = RootModel.dict  # type: ignore
    RootModel.model_json_schema = RootModel.schema  # type: ignore


class SimplePredictBody(BaseModel):
    data: List[Any]
    session_hash: Optional[str] = None


class PredictBody(BaseModel):
    model_config = {"arbitrary_types_allowed": True}

    session_hash: Optional[str] = None
    event_id: Optional[str] = None
    data: List[Any]
    event_data: Optional[Any] = None
    fn_index: Optional[int] = None
    trigger_id: Optional[int] = None
    simple_format: bool = False
    batched: Optional[
        bool
    ] = False  # Whether the data is a batch of samples (i.e. called from the queue if batch=True) or a single sample (i.e. called from the UI)
    request: Optional[
        Request
    ] = None  # dictionary of request headers, query parameters, url, etc. (used to to pass in request for queuing)


class ResetBody(BaseModel):
    event_id: str


class ComponentServerBody(BaseModel):
    session_hash: str
    component_id: int
    fn_name: str
    data: Any


class InterfaceTypes(Enum):
    STANDARD = auto()
    INPUT_ONLY = auto()
    OUTPUT_ONLY = auto()
    UNIFIED = auto()


class GradioBaseModel(ABC):
    def copy_to_dir(self, dir: str | pathlib.Path) -> GradioDataModel:
        if not isinstance(self, (BaseModel, RootModel)):
            raise TypeError("must be used in a Pydantic model")
        dir = pathlib.Path(dir)

        # TODO: Making sure path is unique should be done in caller
        def unique_copy(obj: dict):
            data = FileData(**obj)
            return data._copy_to_dir(
                str(pathlib.Path(dir / secrets.token_hex(10)))
            ).model_dump()

        return self.__class__.from_json(
            x=traverse(
                self.model_dump(),
                unique_copy,
                FileData.is_file_data,
            )
        )

    @classmethod
    @abstractmethod
    def from_json(cls, x) -> GradioDataModel:
        pass


class GradioModel(GradioBaseModel, BaseModel):
    @classmethod
    def from_json(cls, x) -> GradioModel:
        return cls(**x)


class GradioRootModel(GradioBaseModel, RootModel):
    @classmethod
    def from_json(cls, x) -> GradioRootModel:
        return cls(root=x)


GradioDataModel = Union[GradioModel, GradioRootModel]


class FileData(GradioModel):
    path: str  # server filepath
    url: Optional[str] = None  # normalised server url
    size: Optional[int] = None  # size in bytes
    orig_name: Optional[str] = None  # original filename
    mime_type: Optional[str] = None
    is_stream: bool = False

    @property
    def is_none(self):
        return all(
            f is None
            for f in [
                self.path,
                self.url,
                self.size,
                self.orig_name,
                self.mime_type,
            ]
        )

    @classmethod
    def from_path(cls, path: str) -> FileData:
        return cls(path=path)

    def _copy_to_dir(self, dir: str) -> FileData:
        pathlib.Path(dir).mkdir(exist_ok=True)
        new_obj = dict(self)

        if not self.path:
            raise ValueError("Source file path is not set")
        new_name = shutil.copy(self.path, dir)
        new_obj["path"] = new_name
        return self.__class__(**new_obj)

    @classmethod
    def is_file_data(cls, obj: Any):
        if isinstance(obj, dict):
            try:
                return not FileData(**obj).is_none
            except (TypeError, ValidationError):
                return False
        return False


class ListFiles(GradioRootModel):
    root: List[FileData]

    def __getitem__(self, index):
        return self.root[index]

    def __iter__(self):
        return iter(self.root)


@document()
class StaticFiles:
    """
    Static files to be served by the gradio app.

    Static files are not moved to the gradio cache and are served directly from the file system.
    This class is useful when you want to serve files that you know will not be modified during the lifetime of the gradio app (like examples).
    By using StaticFilesm your app will launch faster and it will consume less disk space.

    Example:
        import gradio as gr

        # Paths can be a list of strings or pathlib.Path objects
        # corresponding to filenames or directories.
        gr.StaticFiles(paths=["test/test_files/"])

        # The example files and the default value of the input
        # will not be copied to the gradio cache and will be served directly.
        demo = gr.Interface(
            lambda s: s.rotate(45),
            gr.Image(value="test/test_files/cheetah1.jpg", type="pil"),
            gr.Image(),
            examples=["test/test_files/bus.png"],
        )

        demo.launch()
    """

    all_paths = []

    def __init__(self, paths: list[str | pathlib.Path]) -> None:
        self.paths = paths
        self.all_paths.extend([pathlib.Path(p).resolve() for p in paths])