sift-stack-py 0.3.2__py3-none-any.whl

sift_py/file_attachment/_service_test.py

@@ -0,0 +1,161 @@
+import json
+
+import pytest
+from pytest_mock import MockFixture
+from sift.remote_files.v1.remote_files_pb2 import GetRemoteFileResponse, RemoteFile
+
+from sift_py._internal.test_util.channel import MockChannel
+from sift_py.file_attachment.entity import Entity, EntityType
+from sift_py.file_attachment.metadata import ImageMetadata
+from sift_py.file_attachment.service import FileAttachmentService
+
+
+class MockResponse:
+    status_code: int
+    text: str
+
+    def __init__(self, status_code: int, text: str):
+        self.status_code = status_code
+        self.text = text
+
+    def json(self):
+        return json.loads(self.text)
+
+
+class MockMultipartEncoder:
+    @property
+    def content_type(self):
+        return "multipart/form-data"
+
+
+def test_file_attachments_service_upload_validate_uri():
+    mock_channel = MockChannel()
+
+    svc = FileAttachmentService(
+        mock_channel,
+        {
+            "uri": "https://some_uri.com",
+            "apikey": "123123123",
+        },
+    )
+
+    assert svc is not None
+
+    svc = FileAttachmentService(
+        mock_channel,
+        {
+            "uri": "some_uri.com",
+            "apikey": "123123123",
+        },
+    )
+
+    assert svc is not None
+
+
+def test_file_attachments_service_upload_validate_path(mocker: MockFixture):
+    mock_channel = MockChannel()
+
+    mock_path_is_file = mocker.patch("sift_py.file_attachment._internal.upload.Path.is_file")
+    mock_path_is_file.return_value = False
+
+    with pytest.raises(Exception, match="does not point to a regular file"):
+        svc = FileAttachmentService(
+            mock_channel,
+            {
+                "uri": "some_uri.com",
+                "apikey": "123123123",
+            },
+        )
+
+        svc.upload_attachment(
+            path="some_image.png.gz",
+            entity=Entity(
+                entity_id="123-123-123",
+                entity_type=EntityType.ANNOTATION_LOG,
+            ),
+            metadata=ImageMetadata(
+                width=16,
+                height=9,
+            ),
+        )
+
+
+def test_file_attachments_service_upload_validate_mimetype(mocker: MockFixture):
+    mock_channel = MockChannel()
+
+    mock_path_is_file = mocker.patch("sift_py.file_attachment._internal.upload.Path.is_file")
+    mock_path_is_file.return_value = True
+
+    with pytest.raises(Exception, match="MIME"):
+        svc = FileAttachmentService(
+            mock_channel,
+            {
+                "uri": "some_uri.com",
+                "apikey": "123123123",
+            },
+        )
+
+        svc.upload_attachment(
+            path="some_image.asdlkjfh",
+            entity=Entity(
+                entity_id="123-123-123",
+                entity_type=EntityType.ANNOTATION_LOG,
+            ),
+            metadata=ImageMetadata(
+                width=16,
+                height=9,
+            ),
+        )
+
+
+def test_file_attachments_service_upload_returns_remote_file(mocker: MockFixture):
+    mock_channel = MockChannel()
+
+    mock_path_is_file = mocker.patch("sift_py.file_attachment._internal.upload.Path.is_file")
+    mock_path_is_file.return_value = True
+
+    mocker.patch(
+        "sift_py.file_attachment._internal.upload.open",
+        mocker.mock_open(read_data=b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR"),
+    )
+
+    mock_multipart_encoder = mocker.patch(
+        "sift_py.file_attachment._internal.upload.MultipartEncoder"
+    )
+    mock_multipart_encoder.return_value = MockMultipartEncoder()
+
+    mock_requests_post = mocker.patch("sift_py.file_attachment._internal.upload.requests.post")
+    mock_requests_post.return_value = MockResponse(
+        status_code=200, text=json.dumps({"remoteFile": {"remoteFileId": "abc"}})
+    )
+
+    svc = FileAttachmentService(
+        mock_channel,
+        {
+            "uri": "some_uri.com",
+            "apikey": "123123123",
+        },
+    )
+
+    mock_get_remote_file = mocker.patch.object(
+        svc._remote_file_service_stub,
+        "GetRemoteFile",
+        return_value=GetRemoteFileResponse(remote_file=RemoteFile(remote_file_id="abc")),
+    )
+
+    remote_file = svc.upload_attachment(
+        path="some_image.png.gz",
+        entity=Entity(
+            entity_id="123-123-123",
+            entity_type=EntityType.ANNOTATION_LOG,
+        ),
+        metadata=ImageMetadata(
+            width=16,
+            height=9,
+        ),
+    )
+    mock_get_remote_file.assert_called_once()
+    mock_multipart_encoder.assert_called_once()
+    mock_requests_post.assert_called_once()
+
+    assert remote_file.remote_file_id == "abc"

sift_py/file_attachment/entity.py

@@ -0,0 +1,30 @@
+"""
+Entities represent things that files can be attached to.
+"""
+
+from __future__ import annotations
+
+from enum import Enum
+
+
+class Entity:
+    """
+    An abstract entity that represents the thing that we want to attach files to.
+    """
+
+    entity_id: str
+    entity_type: EntityType
+
+    def __init__(self, entity_id: str, entity_type: EntityType):
+        self.entity_id = entity_id
+        self.entity_type = entity_type
+
+
+class EntityType(Enum):
+    """
+    Represents the types of entities that supports file attachments.
+    """
+
+    RUN = "runs"
+    ANNOTATION = "annotations"
+    ANNOTATION_LOG = "annotation_logs"

sift_py/file_attachment/metadata.py

@@ -0,0 +1,107 @@
+"""
+Module containing optional metadata types to provide to Sift when uploading a file attachment.
+Though optional, providing this information could help improve quality of renders on the Sift app.
+"""
+
+from __future__ import annotations
+
+from datetime import datetime
+from typing import Any, Optional, Type
+
+from google.protobuf.timestamp_pb2 import Timestamp
+from sift.remote_files.v1.remote_files_pb2 import (
+    ImageMetadata as ImageMetadataPb,
+)
+from sift.remote_files.v1.remote_files_pb2 import (
+    VideoMetadata as VideoMetadataPb,
+)
+from typing_extensions import Self
+
+from sift_py._internal.convert.json import AsJson
+from sift_py._internal.convert.protobuf import AsProtobuf
+
+
+class Metadata(AsJson): ...
+
+
+class VideoMetadata(AsProtobuf, Metadata):
+    """
+    Metadata for video media-types i.e. any mimetypes of the following pattern: `video/*`.
+    """
+
+    width: int
+    height: int
+    duration_seconds: float
+    timestamp: Optional[datetime]
+
+    def __init__(
+        self, width: int, height: int, duration_seconds: float, timestamp: Optional[datetime] = None
+    ):
+        self.width = width
+        self.height = height
+        self.duration_seconds = duration_seconds
+        self.timestamp = timestamp
+
+    def as_pb(self, klass: Type[VideoMetadataPb]) -> VideoMetadataPb:
+        if self.timestamp is not None:
+            timestamp_pb = Timestamp()
+            timestamp_pb.FromDatetime(self.timestamp)
+        else:
+            timestamp_pb = None
+
+        return klass(
+            width=self.width,
+            height=self.height,
+            duration_seconds=self.duration_seconds,
+            timestamp=timestamp_pb,
+        )
+
+    @classmethod
+    def from_pb(cls, message: VideoMetadataPb) -> Self:
+        return cls(
+            width=message.width,
+            height=message.height,
+            duration_seconds=message.duration_seconds,
+            timestamp=message.timestamp.ToDateTime(),  # type: ignore
+        )
+
+    def as_json(self) -> Any:
+        timestamp = None if self.timestamp is None else self.timestamp.isoformat()
+        return {
+            "height": self.height,
+            "width": self.width,
+            "duration_seconds": self.duration_seconds,
+            "timestamp": timestamp,
+        }
+
+
+class ImageMetadata(AsProtobuf, Metadata):
+    """
+    Metadata for image media-types i.e. any mimetypes of the following pattern: `image/*`.
+    """
+
+    width: int
+    height: int
+
+    def __init__(self, width: int, height: int):
+        self.width = width
+        self.height = height
+
+    def as_pb(self, klass: Type[ImageMetadataPb]) -> ImageMetadataPb:
+        return klass(
+            width=self.width,
+            height=self.height,
+        )
+
+    @classmethod
+    def from_pb(cls, message: ImageMetadataPb) -> Self:
+        return cls(
+            width=message.width,
+            height=message.height,
+        )
+
+    def as_json(self) -> Any:
+        return {
+            "height": self.height,
+            "width": self.width,
+        }

sift_py/file_attachment/service.py

@@ -0,0 +1,142 @@
+from pathlib import Path
+from typing import List, Optional, Union, cast
+
+from sift.remote_files.v1.remote_files_pb2 import (
+    BatchDeleteRemoteFilesRequest,
+    GetRemoteFileDownloadUrlRequest,
+    GetRemoteFileDownloadUrlResponse,
+    GetRemoteFileRequest,
+    GetRemoteFileResponse,
+    ListRemoteFilesRequest,
+    ListRemoteFilesResponse,
+    RemoteFile,
+)
+from sift.remote_files.v1.remote_files_pb2_grpc import RemoteFileServiceStub
+
+from sift_py.file_attachment._internal.download import download_remote_file
+from sift_py.file_attachment._internal.upload import UploadService
+from sift_py.file_attachment.entity import Entity
+from sift_py.file_attachment.metadata import Metadata
+from sift_py.grpc.transport import SiftChannel
+from sift_py.rest import SiftRestConfig
+
+
+class FileAttachmentService:
+    """
+    Service used to retrieve, upload, download, and delete file attachments. Seee `sift_py.file_attachment`
+    for more information and examples on how to use this service.
+    """
+
+    _remote_file_service_stub: RemoteFileServiceStub
+    _upload_service: UploadService
+
+    def __init__(self, channel: SiftChannel, restconf: SiftRestConfig):
+        self._remote_file_service_stub = RemoteFileServiceStub(channel)
+        self._upload_service = UploadService(restconf)
+
+    def retrieve_attachments(self, entity: Entity) -> List[RemoteFile]:
+        """
+        Retrieves all file attachments for the provided `entity`.
+        """
+
+        filter = f'entity_id=="{entity.entity_id}" && entity_type=="{entity.entity_type.value}"'
+        page_size = 1_000
+        next_page_token = ""
+
+        remote_files: List[RemoteFile] = []
+
+        while True:
+            req = ListRemoteFilesRequest(
+                filter=filter,
+                page_size=page_size,
+                page_token=next_page_token,
+            )
+            res = cast(ListRemoteFilesResponse, self._remote_file_service_stub.ListRemoteFiles(req))
+            remote_files.extend(res.remote_files)
+            next_page_token = res.next_page_token
+
+            if not next_page_token:
+                break
+
+        return remote_files
+
+    def upload_attachment(
+        self,
+        path: Union[str, Path],
+        entity: Entity,
+        metadata: Optional[Metadata],
+        description: Optional[str] = None,
+        organization_id: Optional[str] = None,
+    ) -> RemoteFile:
+        """
+        Uploads a file pointed to by `path` and attaches it to the provided `entity`.
+
+        - `path`: A path to the file to upload to Sift as a file attachment.
+        - `entity`: The entity to attach the file to.
+        - `metadata`: Optional metadata to include with the specific file.
+        - `description`: An optional description to provide for the file attachment.
+        - `organization_id`: Only required if your user belongs to multiple organizations.
+        """
+        remote_file_id = self._upload_service.upload_attachment(
+            path,
+            entity,
+            metadata,
+            description,
+            organization_id,
+        )
+        req = GetRemoteFileRequest(remote_file_id=remote_file_id)
+        res = cast(GetRemoteFileResponse, self._remote_file_service_stub.GetRemoteFile(req))
+        return res.remote_file
+
+    def download_attachment(
+        self,
+        file: Union[RemoteFile, str],
+        out: Optional[Union[str, Path]] = None,
+    ) -> Path:
+        """
+        Downloads a file attachment and saves it locally.
+
+        - `remote_file`: Could either be an instance of `RemoteFile` or the ID of the remote file to download.
+        - `out`: If unspecified, then the file will be downloaded to the current working directory with the original name.
+        """
+
+        if isinstance(file, RemoteFile):
+            remote_file = file
+        else:
+            req = GetRemoteFileRequest(remote_file_id=file)
+            res = cast(GetRemoteFileResponse, self._remote_file_service_stub.GetRemoteFile(req))
+            remote_file = res.remote_file
+
+        output_file_path = (
+            Path(out) if isinstance(out, str) else Path(remote_file.file_name).resolve()
+        )
+
+        download_url_req = GetRemoteFileDownloadUrlRequest(
+            remote_file_id=remote_file.remote_file_id
+        )
+        download_url_res = cast(
+            GetRemoteFileDownloadUrlResponse,
+            self._remote_file_service_stub.GetRemoteFileDownloadUrl(download_url_req),
+        )
+        url = download_url_res.download_url
+
+        download_remote_file(url, output_file_path)
+
+        return output_file_path
+
+    def delete_file_attachments(self, *to_delete: Union[str, RemoteFile]):
+        """
+        Deletes remote files given a set of arguments that could either be instances of `RemoteFile` or the ID
+        of remote files to delete
+        """
+        remote_file_ids = [
+            remote_file.remote_file_id if isinstance(remote_file, RemoteFile) else remote_file
+            for remote_file in to_delete
+        ]
+
+        batch_size = 1_000
+        for i in range(0, len(remote_file_ids), batch_size):
+            batch = remote_file_ids[i : i + batch_size]
+            self._remote_file_service_stub.BatchDeleteRemoteFiles(
+                BatchDeleteRemoteFilesRequest(remote_file_ids=batch)
+            )

sift_py/grpc/__init__.py

@@ -0,0 +1,15 @@
+"""
+This module is primarily concerned with configuring and initializing gRPC connections to the Sift API.
+
+Example of establishing a connection to Sift's gRPC APi:
+
+```python
+from sift_py.grpc.transport import SiftChannelConfig, use_sift_channel
+
+# Be sure not to include the url scheme i.e. 'https://' in the uri.
+sift_channel_config = SiftChannelConfig(uri=SIFT_BASE_URI, apikey=SIFT_API_KEY)
+
+with use_sift_channel(sift_channel_config) as channel:
+    # Connect to Sift
+```
+"""

sift_py/grpc/_async_interceptors/base.py

@@ -0,0 +1,72 @@
+from abc import abstractmethod
+from typing import Any, AsyncIterable, Callable, Iterable, TypeVar, Union
+
+from grpc import aio as grpc_aio
+
+CallType = TypeVar("CallType", bound=grpc_aio.Call)
+Continuation = Callable[[grpc_aio.ClientCallDetails, Any], CallType]
+
+
+class ClientAsyncInterceptor(
+    grpc_aio.UnaryUnaryClientInterceptor,
+    grpc_aio.UnaryStreamClientInterceptor,
+    grpc_aio.StreamUnaryClientInterceptor,
+    grpc_aio.StreamStreamClientInterceptor,
+):
+    @abstractmethod
+    async def intercept(
+        self,
+        method: Callable,
+        request_or_iterator: Any,
+        client_call_details: grpc_aio.ClientCallDetails,
+    ) -> Any:
+        pass
+
+    async def intercept_unary_unary(
+        self,
+        continuation: Continuation[grpc_aio.UnaryUnaryCall],
+        client_call_details: grpc_aio.ClientCallDetails,
+        request: Any,
+    ):
+        return await self.intercept(_async_swap_args(continuation), request, client_call_details)
+
+    async def intercept_unary_stream(
+        self,
+        continuation: Continuation[grpc_aio.UnaryStreamCall],
+        client_call_details: grpc_aio.ClientCallDetails,
+        request: Any,
+    ):
+        return await self.intercept(_async_swap_args(continuation), request, client_call_details)
+
+    async def intercept_stream_unary(
+        self,
+        continuation: Continuation[grpc_aio.StreamUnaryCall],
+        client_call_details: grpc_aio.ClientCallDetails,
+        request_iterator: Union[Iterable[Any], AsyncIterable[Any]],
+    ):
+        return await self.intercept(
+            _async_swap_args(continuation), request_iterator, client_call_details
+        )
+
+    async def intercept_stream_stream(
+        self,
+        continuation: Continuation[grpc_aio.StreamStreamCall],
+        client_call_details: grpc_aio.ClientCallDetails,
+        request_iterator: Union[Iterable[Any], AsyncIterable[Any]],
+    ):
+        return await self.intercept(
+            _async_swap_args(continuation), request_iterator, client_call_details
+        )
+
+
+def _async_swap_args(fn: Callable[[Any, Any], Any]) -> Callable[[Any, Any], Any]:
+    """
+    Continuations are typed in such a way that details are the first argument, and the request second.
+    Code generated from protobuf however takes in the request first, then the details. Weird grpc library
+    quirk. This utility just flips the arguments.
+    """
+
+    async def new_fn(x, y):
+        return await fn(y, x)
+
+    return new_fn

sift_py/grpc/_async_interceptors/metadata.py

@@ -0,0 +1,36 @@
+from __future__ import annotations
+
+from typing import Any, Callable, List, Tuple, cast
+
+from grpc import aio as grpc_aio
+
+from sift_py.grpc._async_interceptors.base import ClientAsyncInterceptor
+
+Metadata = List[Tuple[str, str]]
+
+
+class MetadataAsyncInterceptor(ClientAsyncInterceptor):
+    metadata: Metadata
+
+    """
+    Interceptor to add metadata to all async unary and streaming RPCs
+    """
+
+    def __init__(self, metadata: Metadata):
+        self.metadata = metadata
+
+    async def intercept(
+        self,
+        method: Callable,
+        request_or_iterator: Any,
+        client_call_details: grpc_aio.ClientCallDetails,
+    ):
+        call_details = cast(grpc_aio.ClientCallDetails, client_call_details)
+        new_details = grpc_aio.ClientCallDetails(
+            call_details.method,
+            call_details.timeout,
+            self.metadata,
+            call_details.credentials,
+            call_details.wait_for_ready,
+        )
+        return await method(request_or_iterator, new_details)

sift_py/grpc/_interceptors/base.py

@@ -0,0 +1,61 @@
+from abc import abstractmethod
+from typing import Any, Callable, Iterator
+
+import grpc
+
+Continuation = Callable[[grpc.ClientCallDetails, Any], Any]
+
+
+class ClientInterceptor(
+    grpc.StreamStreamClientInterceptor,
+    grpc.StreamUnaryClientInterceptor,
+    grpc.UnaryStreamClientInterceptor,
+    grpc.UnaryUnaryClientInterceptor,
+):
+    @abstractmethod
+    def intercept(
+        self,
+        method: Continuation,
+        request_or_iterator: Any,
+        client_call_details: grpc.ClientCallDetails,
+    ):
+        pass
+
+    def intercept_unary_unary(
+        self,
+        continuation: Continuation,
+        client_call_details: grpc.ClientCallDetails,
+        request: Any,
+    ):
+        return self.intercept(_swap_args(continuation), request, client_call_details)
+
+    def intercept_stream_unary(
+        self,
+        continuation: Continuation,
+        client_call_details: grpc.ClientCallDetails,
+        request_iterator: Iterator[Any],
+    ):
+        return self.intercept(_swap_args(continuation), request_iterator, client_call_details)
+
+    def intercept_unary_stream(
+        self,
+        continuation: Continuation,
+        client_call_details: grpc.ClientCallDetails,
+        request: Any,
+    ):
+        return self.intercept(_swap_args(continuation), request, client_call_details)
+
+    def intercept_stream_stream(
+        self,
+        continuation: Continuation,
+        client_call_details: grpc.ClientCallDetails,
+        request_iterator: Iterator[Any],
+    ):
+        return self.intercept(_swap_args(continuation), request_iterator, client_call_details)
+
+
+def _swap_args(fn: Callable[[Any, Any], Any]) -> Callable[[Any, Any], Any]:
+    def new_fn(x, y):
+        return fn(y, x)
+
+    return new_fn

sift_py/grpc/_interceptors/context.py

@@ -0,0 +1,25 @@
+from typing import Optional, Sequence, Tuple, Union
+
+import grpc
+
+
+class ClientCallDetails(grpc.ClientCallDetails):
+    method: str
+    timeout: Optional[float]
+    metadata: Optional[Sequence[Tuple[str, Union[str, bytes]]]]
+    credentials: Optional[grpc.CallCredentials]
+    wait_for_ready: Optional[bool]
+
+    def __init__(
+        self,
+        method: str,
+        timeout: Optional[float],
+        metadata: Optional[Sequence[Tuple[str, Union[str, bytes]]]],
+        credentials: Optional[grpc.CallCredentials],
+        wait_for_ready: Optional[bool],
+    ):
+        self.method = method
+        self.timeout = timeout
+        self.metadata = metadata
+        self.credentials = credentials
+        self.wait_for_ready = wait_for_ready

sift_py/grpc/_interceptors/metadata.py

@@ -0,0 +1,33 @@
+from typing import Any, List, Tuple, cast
+
+import grpc
+
+from sift_py.grpc._interceptors.base import ClientInterceptor, Continuation
+from sift_py.grpc._interceptors.context import ClientCallDetails
+
+Metadata = List[Tuple[str, str]]
+
+
+class MetadataInterceptor(ClientInterceptor):
+    metadata: Metadata
+
+    def __init__(self, metadata: Metadata):
+        self.metadata = metadata
+
+    def intercept(
+        self,
+        method: Continuation,
+        request_or_iterator: Any,
+        client_call_details: grpc.ClientCallDetails,
+    ):
+        details = cast(ClientCallDetails, client_call_details)
+
+        new_details = ClientCallDetails(
+            method=details.method,
+            timeout=details.timeout,
+            credentials=details.credentials,
+            wait_for_ready=details.wait_for_ready,
+            metadata=self.metadata,
+        )
+
+        return method(request_or_iterator, new_details)