Source code for ai2_kit.core.artifact

from .pydantic import BaseModel
from typing import Optional, Mapping, TypedDict
import os
import copy


def __ArtifactDict():
    class ArtifactDict(TypedDict):
        """
        A dict representation of Artifact.
        Use this when you need to run a remote function call as pydantic model is not pickleable.

        referrer is not included in this dict as it is not pickleable.
        """
        url: str
        attrs: dict
        executor: Optional[str]
        format: Optional[str]
        includes: Optional[str]
        key: Optional[str]
    return ArtifactDict
ArtifactDict = __ArtifactDict()


class Artifact(BaseModel):

    key: Optional[str]
    executor: Optional[str]
    url: str
    format: Optional[str]
    includes: Optional[str]
    attrs: dict = dict()

    @classmethod
    def of(cls,
           url: str,
           key: Optional[str] = None,
           executor: Optional[str] = None,
           includes: Optional[str] = None,
           attrs: Optional[dict] = None,
           format: Optional[str] = None,):
        """Create an Artifact instance. Use this instead of __init__ to avoid type error of pydantic"""
        return cls(url=url,
                   executor=executor,
                   format=format,
                   includes=includes,
                   attrs=dict() if attrs is None else copy.deepcopy(attrs),  # deepcopy to reduce chance of mistake
                   key=key)

    def to_dict(self) -> ArtifactDict:
        """Convert to a dict representation.
        Use this when you need to run a remote function call as pydantic model is not pickleable."""
        return self.dict() # type: ignore

    def join(self, *paths, **kwargs):
        url = os.path.join(self.url, *paths)
        return Artifact.of(url=url, executor=self.executor, **kwargs)


ArtifactMap = Mapping[str, Artifact]