Source code for mlflow.entities.file_info

from mlflow.entities._mlflow_object import _MlflowObject
from mlflow.protos.service_pb2 import FileInfo as ProtoFileInfo


class FileInfo(_MlflowObject):
    """
    Metadata about a file or directory.
    """

    def __init__(self, path: str, is_dir: bool, file_size: int | None) -> None:
        self._path = path
        self._is_dir = is_dir
        self._bytes = file_size

    def __eq__(self, other: object) -> bool:
        if type(other) is type(self):
            return self.__dict__ == other.__dict__
        return False

    @property
    def path(self) -> str:
        """String path of the file or directory."""
        return self._path

    @property
    def is_dir(self) -> bool:
        """Whether the FileInfo corresponds to a directory."""
        return self._is_dir

    @property
    def file_size(self) -> int | None:
        """Size of the file or directory. If the FileInfo is a directory, returns None."""
        return self._bytes

    def to_proto(self) -> ProtoFileInfo:
        proto = ProtoFileInfo()
        proto.path = self.path
        proto.is_dir = self.is_dir
        if self.file_size:
            proto.file_size = self.file_size
        return proto

    @classmethod
    def from_proto(cls, proto: ProtoFileInfo) -> "FileInfo":
        return cls(proto.path, proto.is_dir, proto.file_size)