Skip to content

Commit

Permalink
Merge pull request #4 from tektronix/artifacts
Browse files Browse the repository at this point in the history
File artifacts
  • Loading branch information
TJBIII authored Jan 11, 2022
2 parents cb278bf + d333b1a commit 410bbec
Show file tree
Hide file tree
Showing 17 changed files with 414 additions and 56 deletions.
1 change: 1 addition & 0 deletions docs/reference/models.rst
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ Models
:maxdepth: 2
:caption: Models

models/artifact
models/file
models/folder
models/member
Expand Down
8 changes: 8 additions & 0 deletions docs/reference/models/artifact.rst
Original file line number Diff line number Diff line change
@@ -0,0 +1,8 @@
.. _artifact:

Artifact
========

.. autoclass:: tekdrive.models.Artifact
:inherited-members:
:exclude-members: parse
2 changes: 2 additions & 0 deletions tekdrive/enums.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ class SharingType(Enum):


class ObjectType(Enum):
ARTIFACT = "ARTIFACT"
FILE = "FILE"
FOLDER = "FOLDER"

Expand All @@ -19,6 +20,7 @@ class FolderType(Enum):


class ErrorCode(Enum):
ARTIFACT_NOT_FOUND = "ARTIFACT_NOT_FOUND"
FILE_GONE = "FILE_GONE"
FILE_NOT_FOUND = "FILE_NOT_FOUND"
FOLDER_GONE = "FOLDER_GONE"
Expand Down
4 changes: 4 additions & 0 deletions tekdrive/exceptions.py
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,10 @@ def request_id(self):
return self.headers.get("X-Request-Id")


class ArtifactNotFoundAPIException(TekDriveAPIException):
"""Indicate artifact is not found."""


class FileNotFoundAPIException(TekDriveAPIException):
"""Indicate file is not found."""

Expand Down
1 change: 1 addition & 0 deletions tekdrive/models/__init__.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
from .drive.file import File # noqa
from .drive.folder import Folder # noqa
from .drive.artifact import Artifact, ArtifactsList # noqa
from .drive.member import Member, MembersList # noqa
from .drive.plan import Plan # noqa
from .drive.trash import Trash # noqa
Expand Down
59 changes: 59 additions & 0 deletions tekdrive/models/drive/artifact.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,59 @@
"""Provides the Artifact class."""
from typing import TYPE_CHECKING, Optional, Dict, Any, Union

from .base import DriveBase, Downloadable
from ..base import BaseList
from ...routing import Route, ENDPOINTS

if TYPE_CHECKING:
from .. import TekDrive


class Artifact(DriveBase, Downloadable):
"""
Represents a file artifact.
Attributes:
bytes (str): Artifact size in bytes.
created_at (datetime): When the artifact was created.
context_type (str): Context to identify the type of artifact such as ``"SETTING"`` or ``"CHANNEL"``.
file_id (str): ID of the file that the artifact is associated with.
file_type (str): Artifact file type such as ``"TSS"`` or ``"SET"``.
id (str): Unique artifact ID.
name (str): Artifact name.
parent_artifact_id: ID of the parent artifact.
updated_at (datetime, optional): When the artifact was last updated.
"""

STR_FIELD = "id"

@classmethod
def from_data(cls, tekdrive, data):
return cls(tekdrive, data)

def __init__(
self,
tekdrive: "TekDrive",
_data: Optional[Dict[str, Any]] = None,
):
super().__init__(tekdrive, _data=_data)

def __setattr__(
self,
attribute: str,
value: Union[str, int, Dict[str, Any]],
):
super().__setattr__(attribute, value)

def _fetch_download_url(self):
route = Route("GET", ENDPOINTS["file_artifact_download"], file_id=self.file_id, artifact_id=self.id)
download_details = self._tekdrive.request(route)
return download_details["download_url"]


class ArtifactsList(BaseList):
"""List of artifacts"""

_parent = None

LIST_ATTRIBUTE = "artifacts"
64 changes: 63 additions & 1 deletion tekdrive/models/drive/base.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,11 @@
"""Provide the DriveBase class."""
from typing import TYPE_CHECKING, Any, Dict, Optional, Union
import os
import requests
from abc import ABC, abstractmethod
from typing import TYPE_CHECKING, Any, Dict, IO, Optional, Union

from ..base import TekDriveBase
from ...exceptions import ClientException, TekDriveStorageException

if TYPE_CHECKING:
from ... import TekDrive
Expand Down Expand Up @@ -70,3 +74,61 @@ def _reset_attributes(self, *attributes):
if attribute in self.__dict__:
del self.__dict__[attribute]
self._fetched = False


class Downloadable(ABC):
"""Abstract base class for download functionality."""

@abstractmethod
def _fetch_download_url(self):
pass

def _download_from_storage(self):
download_url = self._fetch_download_url()
try:
r = requests.get(
download_url,
)
r.raise_for_status()
return r.content
except requests.exceptions.HTTPError as exception:
raise TekDriveStorageException("Download failed") from exception

def download(self, path_or_writable: Union[str, IO] = None) -> None:
"""
Download contents.
Args:
path_or_writable: Path to a local file or a writable stream
where contents will be written.
Raises:
ClientException: If invalid file path is given.
Examples:
Download to local file using path::
here = os.path.dirname(__file__)
contents_path = os.path.join(here, "test_file_overwrite.txt")
file.download(contents_path)
Download using writable stream::
with open("./download.csv", "wb") as f:
file.download(f)
"""
if path_or_writable is None:
# return content directly
return self._download_from_storage()

if isinstance(path_or_writable, str):
file_path = path_or_writable

if not os.path.exists(file_path):
raise ClientException(f"File '{file_path}' does not exist.")

with open(file_path, "wb") as file:
file.write(self._download_from_storage())
else:
writable = path_or_writable
writable.write(self._download_from_storage())
103 changes: 50 additions & 53 deletions tekdrive/models/drive/file.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,8 @@
from ...routing import Route, ENDPOINTS
from ...exceptions import ClientException, TekDriveStorageException
from ...utils.casing import to_snake_case, to_camel_case
from .base import DriveBase
from .base import DriveBase, Downloadable
from .artifact import Artifact, ArtifactsList
from .member import Member, MembersList
from .user import PartialUser
from ...enums import ObjectType
Expand All @@ -17,7 +18,7 @@
from .. import TekDrive


class File(DriveBase):
class File(DriveBase, Downloadable):
"""
A class representing a TekDrive file.
Expand Down Expand Up @@ -120,17 +121,6 @@ def _upload_to_storage(self, file: IO):
except requests.exceptions.HTTPError as exception:
raise TekDriveStorageException("Upload failed") from exception

def _download_from_storage(self):
download_url = self._fetch_download_url()
try:
r = requests.get(
download_url,
)
r.raise_for_status()
return r.content
except requests.exceptions.HTTPError as exception:
raise TekDriveStorageException("Upload failed") from exception

@staticmethod
def _create(
_tekdrive,
Expand All @@ -154,6 +144,52 @@ def _create(

return new_file

def artifacts(self, flat: bool = False) -> ArtifactsList:
"""
Get a list of file artifacts.
Args:
flat: Return artifacts in flat list with no child nesting?
Examples:
Iterate over all file artifacts::
for artifact in file.artifacts():
print(artifact.name)
Returns:
List [ :ref:`artifact` ]
"""
params = to_camel_case(dict(flat=flat))

route = Route("GET", ENDPOINTS["file_artifacts"], file_id=self.id)
artifacts = self._tekdrive.request(route, params=params)
artifacts._parent = self
return artifacts

def artifact(self, artifact_id: str, depth: int = 1) -> Artifact:
"""
Get a file artifact by ID.
Args:
artifact_id: Unique ID for the artifact
depth: How many nested levels of child artifacts to return.
Examples:
Load artifact with children up to 3 levels deep::
artifact_id = "017820c4-03ba-4e9d-be2f-e0ba346ddd9b"
artifact = file.artifact(artifact_id, depth=3)
Returns:
:ref:`artifact`
"""
params = to_camel_case(dict(depth=depth))

route = Route("GET", ENDPOINTS["file_artifact"], file_id=self.id, artifact_id=artifact_id)
artifact = self._tekdrive.request(route, params=params)
return artifact

def members(self) -> MembersList:
"""
Get a list of file members.
Expand Down Expand Up @@ -225,7 +261,7 @@ def upload(self, path_or_readable: Union[str, IO]) -> None:
Upload using readable stream::
with open("./test_file.txt"), "rb") as f:
with open("./test_file.txt", "rb") as f:
new_file.upload(f)
"""
# TODO: multipart upload support
Expand All @@ -241,45 +277,6 @@ def upload(self, path_or_readable: Union[str, IO]) -> None:
readable = path_or_readable
self._upload_to_storage(readable)

def download(self, path_or_writable: Union[str, IO] = None) -> None:
"""
Download file contents.
Args:
path_or_writable: Path to a local file or a writable stream
where file contents will be written.
Raises:
ClientException: If invalid file path is given.
Examples:
Download to local file using path::
here = os.path.dirname(__file__)
contents_path = os.path.join(here, "test_file_overwrite.txt")
file.download(contents_path)
Download using writable stream::
with open("./download.csv"), "wb") as f:
file.download(f)
"""
if path_or_writable is None:
# return content directly
return self._download_from_storage()

if isinstance(path_or_writable, str):
file_path = path_or_writable

if not os.path.exists(file_path):
raise ClientException(f"File '{file_path}' does not exist.")

with open(file_path, "wb") as file:
file.write(self._download_from_storage())
else:
writable = path_or_writable
writable.write(self._download_from_storage())

def move(self, parent_folder_id: str) -> None:
"""
Move file to a different folder.
Expand Down
10 changes: 10 additions & 0 deletions tekdrive/models/parser.py
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,12 @@ def __init__(self, tekdrive: "TekDrive", models: Optional[Dict[str, Any]] = None
self._tekdrive = tekdrive
self.models = {} if models is None else models

def _is_artifact(self, data: dict) -> bool:
return data.get("type") == ObjectType.ARTIFACT.value

def _is_artifacts_list(self, data: dict) -> bool:
return "artifacts" in data

def _is_file(self, data: dict) -> bool:
return data.get("type") == ObjectType.FILE.value

Expand Down Expand Up @@ -73,6 +79,10 @@ def _parse_dict(self, data: dict):
model = self.models["File"]
elif self._is_folder(data):
model = self.models["Folder"]
elif self._is_artifacts_list(data):
model = self.models["ArtifactsList"]
elif self._is_artifact(data):
model = self.models["Artifact"]
elif self._is_members_list(data):
model = self.models["MembersList"]
elif self._is_member(data):
Expand Down
3 changes: 3 additions & 0 deletions tekdrive/routing.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,9 @@ def __init__(self, method, path_template, **params):


ENDPOINTS = {
"file_artifacts": "/file/{file_id}/artifacts",
"file_artifact": "/file/{file_id}/artifacts/{artifact_id}",
"file_artifact_download": "/file/{file_id}/artifacts/{artifact_id}/contents",
"file_create": "/file",
"file_details": "/file/{file_id}",
"file_delete": "/file/{file_id}",
Expand Down
2 changes: 1 addition & 1 deletion tekdrive/settings.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
"""Settings/constants"""

__version__ = "1.0.0"
__version__ = "1.1.0"

RATELIMIT_SECONDS = 1
TIMEOUT = 15
Expand Down
2 changes: 2 additions & 0 deletions tekdrive/tekdrive.py
Original file line number Diff line number Diff line change
Expand Up @@ -64,6 +64,8 @@ def __exit__(self, *_args):

def _create_model_map(self):
model_map = {
"Artifact": models.Artifact,
"ArtifactsList": models.ArtifactsList,
"File": models.File,
"Folder": models.Folder,
"Member": models.Member,
Expand Down
Loading

0 comments on commit 410bbec

Please sign in to comment.