Skip to content

Commit 59d9756

Browse files
haakonvthaakonvt
authored
feat(Data Modeling, Files): Extend DM Files with all download methods (DM-3842) (#2653)
1 parent fcd74e6 commit 59d9756

5 files changed

Lines changed: 347 additions & 38 deletions

File tree

‎cognite/client/_api/data_modeling/files.py‎

Lines changed: 111 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
from __future__ import annotations
22

33
from collections.abc import Sequence
4+
from pathlib import Path
45
from typing import TYPE_CHECKING, Any, NoReturn, overload
56

67
from cognite.client._api_client import APIClient
@@ -44,17 +45,119 @@ def __init__(self, config: ClientConfig, api_version: str | None, cognite_client
4445
super().__init__(config, api_version, cognite_client)
4546
self._files_api = cognite_client.files
4647

47-
async def retrieve_download_urls(self, *args: Any, **kwargs: Any) -> NoReturn:
48-
raise NotImplementedError("This method is not implemented yet!")
48+
async def retrieve_download_urls(
49+
self,
50+
node_ids: NodeId | tuple[str, str] | Sequence[NodeId | tuple[str, str]],
51+
*,
52+
extended_expiration: bool = False,
53+
) -> dict[NodeId, str]:
54+
"""`Get download URLs for one or more files by instance ID. <https://api-docs.cognite.com/20230101/tag/Files/operation/downloadLinks>`_
4955
50-
async def download(self, *args: Any, **kwargs: Any) -> NoReturn:
51-
raise NotImplementedError("This method is not implemented yet!")
56+
Args:
57+
node_ids (NodeId | tuple[str, str] | Sequence[NodeId | tuple[str, str]]): Instance ID or list of instance IDs.
58+
extended_expiration (bool): Extend expiration time of download URL to 1 hour. Defaults to False.
5259
53-
async def download_to_path(self, *args: Any, **kwargs: Any) -> NoReturn:
54-
raise NotImplementedError("This method is not implemented yet!")
60+
Returns:
61+
dict[NodeId, str]: Dictionary mapping each instance ID to its download URL.
5562
56-
async def download_bytes(self, *args: Any, **kwargs: Any) -> NoReturn:
57-
raise NotImplementedError("This method is not implemented yet!")
63+
Examples:
64+
65+
Get a download URL for a single file:
66+
67+
>>> from cognite.client import CogniteClient
68+
>>> from cognite.client.data_classes.data_modeling import NodeId
69+
>>> client = CogniteClient()
70+
>>> urls = client.data_modeling.files.retrieve_download_urls(
71+
... NodeId("my-space", "my-file")
72+
... )
73+
74+
Get download URLs for multiple files:
75+
76+
>>> urls = client.data_modeling.files.retrieve_download_urls(
77+
... [("my-space", "file-1"), ("my-space", "file-2")]
78+
... )
79+
"""
80+
return await self._files_api.retrieve_download_urls(
81+
instance_id=node_ids,
82+
extended_expiration=extended_expiration,
83+
)
84+
85+
async def download(
86+
self,
87+
directory: Path,
88+
node_ids: NodeId | tuple[str, str] | Sequence[NodeId | tuple[str, str]],
89+
keep_directory_structure: bool = False,
90+
resolve_duplicate_file_names: bool = False,
91+
) -> None:
92+
"""`Download files by instance ID. <https://api-docs.cognite.com/20230101/tag/Files/operation/downloadLinks>`_
93+
94+
Streams all files to disk, never keeping more than 2MB in memory per worker.
95+
96+
Args:
97+
directory (Path): Directory to download the file(s) to.
98+
node_ids (NodeId | tuple[str, str] | Sequence[NodeId | tuple[str, str]]): Instance ID or list of instance IDs.
99+
keep_directory_structure (bool): Whether to keep the directory hierarchy from CDF, creating subdirectories as needed.
100+
resolve_duplicate_file_names (bool): Whether to resolve duplicate file names by appending a number.
101+
102+
Examples:
103+
104+
Download files by instance ID into directory 'my_directory':
105+
106+
>>> from pathlib import Path
107+
>>> from cognite.client import CogniteClient
108+
>>> from cognite.client.data_classes.data_modeling import NodeId
109+
>>> client = CogniteClient()
110+
>>> client.data_modeling.files.download(
111+
... directory=Path("my_directory"),
112+
... node_ids=[NodeId("my-space", "file-1"), NodeId("my-space", "file-2")],
113+
... )
114+
"""
115+
await self._files_api.download(
116+
directory=directory,
117+
instance_id=node_ids,
118+
keep_directory_structure=keep_directory_structure,
119+
resolve_duplicate_file_names=resolve_duplicate_file_names,
120+
)
121+
122+
async def download_to_path(self, path: Path, node_id: NodeId | tuple[str, str]) -> None:
123+
"""Download a file to a specific path by instance ID.
124+
125+
Args:
126+
path (Path): Download to this path.
127+
node_id (NodeId | tuple[str, str]): Instance ID of the file to download.
128+
129+
Examples:
130+
131+
Download a file by instance ID:
132+
133+
>>> from cognite.client import CogniteClient
134+
>>> from cognite.client.data_classes.data_modeling import NodeId
135+
>>> client = CogniteClient()
136+
>>> client.data_modeling.files.download_to_path(
137+
... Path("~/mydir/my_file.txt"), NodeId("my-space", "my-file")
138+
... )
139+
"""
140+
await self._files_api.download_to_path(path=path, instance_id=node_id)
141+
142+
async def download_bytes(self, node_id: NodeId | tuple[str, str]) -> bytes:
143+
"""Download a file as bytes by instance ID.
144+
145+
Args:
146+
node_id (NodeId | tuple[str, str]): Instance ID of the file.
147+
148+
Returns:
149+
bytes: The file content.
150+
151+
Examples:
152+
153+
Download a file's content into memory by instance ID:
154+
155+
>>> from cognite.client import CogniteClient
156+
>>> from cognite.client.data_classes.data_modeling import NodeId
157+
>>> client = CogniteClient()
158+
>>> content = client.data_modeling.files.download_bytes(NodeId("my-space", "my-file"))
159+
"""
160+
return await self._files_api.download_bytes(instance_id=node_id)
58161

59162
async def upload(self, *args: Any, **kwargs: Any) -> NoReturn:
60163
raise NotImplementedError("This method is not implemented yet!")

‎cognite/client/_api/files.py‎

Lines changed: 53 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -1104,23 +1104,61 @@ async def _complete_multipart_upload(self, session: FileMultipartUploadSession)
11041104
semaphore=self._get_semaphore("upload"),
11051105
)
11061106

1107+
@overload
1108+
async def retrieve_download_urls(
1109+
self,
1110+
id: int | Sequence[int],
1111+
external_id: None = None,
1112+
instance_id: None = None,
1113+
extended_expiration: bool = False,
1114+
) -> dict[int, str]: ...
1115+
1116+
@overload
1117+
async def retrieve_download_urls(
1118+
self,
1119+
id: None = None,
1120+
*,
1121+
external_id: str | SequenceNotStr[str],
1122+
instance_id: None = None,
1123+
extended_expiration: bool = False,
1124+
) -> dict[str, str]: ...
1125+
1126+
@overload
1127+
async def retrieve_download_urls(
1128+
self,
1129+
id: None = None,
1130+
external_id: None = None,
1131+
*,
1132+
instance_id: NodeId | tuple[str, str] | Sequence[NodeId | tuple[str, str]],
1133+
extended_expiration: bool = False,
1134+
) -> dict[NodeId, str]: ...
1135+
1136+
@overload
11071137
async def retrieve_download_urls(
11081138
self,
11091139
id: int | Sequence[int] | None = None,
11101140
external_id: str | SequenceNotStr[str] | None = None,
1111-
instance_id: NodeId | Sequence[NodeId] | None = None,
1141+
instance_id: NodeId | tuple[str, str] | Sequence[NodeId | tuple[str, str]] | None = None,
11121142
extended_expiration: bool = False,
1113-
) -> dict[int | str | NodeId, str]:
1143+
) -> dict[int | str | NodeId, str]: ...
1144+
1145+
async def retrieve_download_urls(
1146+
self,
1147+
id: int | Sequence[int] | None = None,
1148+
external_id: str | SequenceNotStr[str] | None = None,
1149+
instance_id: NodeId | tuple[str, str] | Sequence[NodeId | tuple[str, str]] | None = None,
1150+
extended_expiration: bool = False,
1151+
) -> dict[int, str] | dict[str, str] | dict[NodeId, str] | dict[int | str | NodeId, str]:
11141152
"""Get download links by id or external id.
11151153
11161154
Args:
11171155
id (int | Sequence[int] | None): Id or list of ids.
11181156
external_id (str | SequenceNotStr[str] | None): External id or list of external ids.
1119-
instance_id (NodeId | Sequence[NodeId] | None): Instance id or list of instance ids.
1157+
instance_id (NodeId | tuple[str, str] | Sequence[NodeId | tuple[str, str]] | None): Instance id or list of instance ids.
11201158
extended_expiration (bool): Extend expiration time of download url to 1 hour. Defaults to false.
11211159
11221160
Returns:
1123-
dict[int | str | NodeId, str]: Dictionary containing download urls.
1161+
dict[int, str] | dict[str, str] | dict[NodeId, str] | dict[int | str | NodeId, str]: Dictionary containing download urls.
11241162
"""
11251163
identifiers = IdentifierSequence.load(ids=id, external_ids=external_id, instance_ids=instance_id)
11261164

@@ -1176,7 +1214,7 @@ async def download(
11761214
directory: str | Path,
11771215
id: int | Sequence[int] | None = None,
11781216
external_id: str | SequenceNotStr[str] | None = None,
1179-
instance_id: NodeId | Sequence[NodeId] | None = None,
1217+
instance_id: NodeId | tuple[str, str] | Sequence[NodeId | tuple[str, str]] | None = None,
11801218
keep_directory_structure: bool = False,
11811219
resolve_duplicate_file_names: bool = False,
11821220
) -> None:
@@ -1197,9 +1235,8 @@ async def download(
11971235
directory (str | Path): Directory to download the file(s) to.
11981236
id (int | Sequence[int] | None): Id or list of ids
11991237
external_id (str | SequenceNotStr[str] | None): External ID or list of external ids.
1200-
instance_id (NodeId | Sequence[NodeId] | None): Instance ID or list of instance ids.
1201-
keep_directory_structure (bool): Whether or not to keep the directory hierarchy in CDF,
1202-
creating subdirectories as needed below the given directory.
1238+
instance_id (NodeId | tuple[str, str] | Sequence[NodeId | tuple[str, str]] | None): Instance ID or list of instance ids.
1239+
keep_directory_structure (bool): Whether or not to keep the directory hierarchy in CDF, creating subdirectories as needed below the given directory.
12031240
resolve_duplicate_file_names (bool): Whether or not to resolve duplicate file names by appending a number on duplicate file names
12041241
12051242
Examples:
@@ -1335,15 +1372,19 @@ async def _download_file_to_path(self, download_link: str, path: Path) -> None:
13351372
file.write(chunk)
13361373

13371374
async def download_to_path(
1338-
self, path: Path | str, id: int | None = None, external_id: str | None = None, instance_id: NodeId | None = None
1375+
self,
1376+
path: Path | str,
1377+
id: int | None = None,
1378+
external_id: str | None = None,
1379+
instance_id: NodeId | tuple[str, str] | None = None,
13391380
) -> None:
13401381
"""Download a file to a specific target.
13411382
13421383
Args:
13431384
path (Path | str): Download to this path.
13441385
id (int | None): Id of of the file to download.
13451386
external_id (str | None): External id of the file to download.
1346-
instance_id (NodeId | None): Instance id of the file to download.
1387+
instance_id (NodeId | tuple[str, str] | None): Instance id of the file to download.
13471388
13481389
Examples:
13491390
@@ -1362,14 +1403,14 @@ async def download_to_path(
13621403
await self._download_file_to_path(download_link, path)
13631404

13641405
async def download_bytes(
1365-
self, id: int | None = None, external_id: str | None = None, instance_id: NodeId | None = None
1406+
self, id: int | None = None, external_id: str | None = None, instance_id: NodeId | tuple[str, str] | None = None
13661407
) -> bytes:
13671408
"""Download a file as bytes.
13681409
13691410
Args:
13701411
id (int | None): Id of the file
13711412
external_id (str | None): External id of the file
1372-
instance_id (NodeId | None): Instance id of the file
1413+
instance_id (NodeId | tuple[str, str] | None): Instance id of the file
13731414
13741415
Examples:
13751416

0 commit comments

Comments
 (0)