diff --git a/tensorbay/client/segment.py b/tensorbay/client/segment.py index 626c92943..c6698dfc4 100644 --- a/tensorbay/client/segment.py +++ b/tensorbay/client/segment.py @@ -123,6 +123,19 @@ def _list_urls(self, offset: int = 0, limit: int = 128) -> Dict[str, Any]: response = self._client.open_api_do("GET", "data/urls", self._dataset_id, params=params) return response.json() # type: ignore[no-any-return] + def _get_data_details(self, remote_path: str) -> Dict[str, Any]: + params: Dict[str, Any] = { + "segmentName": self._name, + "remotePath": remote_path, + } + params.update(self._status.get_status_info()) + + if config.is_internal: + params["isInternal"] = True + + response = self._client.open_api_do("GET", "data/details", self._dataset_id, params=params) + return response.json() # type: ignore[no-any-return] + def _list_data_details(self, offset: int = 0, limit: int = 128) -> Dict[str, Any]: params: Dict[str, Any] = { "segmentName": self._name, @@ -137,6 +150,20 @@ def _list_data_details(self, offset: int = 0, limit: int = 128) -> Dict[str, Any response = self._client.open_api_do("GET", "data/details", self._dataset_id, params=params) return response.json() # type: ignore[no-any-return] + def _get_mask_url(self, mask_type: str, remote_path: str) -> str: + params: Dict[str, Any] = { + "segmentName": self._name, + "maskType": mask_type, + "remotePath": remote_path, + } + params.update(self._status.get_status_info()) + + if config.is_internal: + params["isInternal"] = True + + response = self._client.open_api_do("GET", "masks/urls", self._dataset_id, params=params) + return response.json()["urls"][0]["url"] # type: ignore[no-any-return] + def _list_mask_urls(self, mask_type: str, offset: int = 0, limit: int = 128) -> Dict[str, Any]: params: Dict[str, Any] = { "segmentName": self._name, @@ -653,6 +680,45 @@ def list_data_paths(self) -> PagingList[str]: """ return PagingList(self._generate_data_paths, 128) + def get_data(self, remote_path: str) -> RemoteData: + """Get required Data object from a dataset segment. + + Arguments: + remote_path: The remote paths of the required data. + + Returns: + :class:`~tensorbay.dataset.data.RemoteData`. + + """ + response = self._get_data_details(remote_path) + data_details = response["dataDetails"][0] + + mask_urls = {} + for key in _MASK_KEYS: + mask_urls[key] = self._get_mask_url(key.upper(), remote_path) + + data = RemoteData.from_response_body( + data_details, + url=URL(data_details["url"], lambda: self._get_url(remote_path)), + cache_path=self._cache_path, + ) + label = data.label + for key in _MASK_KEYS: + mask = getattr(label, key, None) + if mask: + # pylint: disable=protected-access + mask.url = ( + URL( + mask_urls[key], + lambda k=key.upper(), r=remote_path: ( # type: ignore[misc] + self._get_mask_url(k, r) + ), + ), + ) + mask.cache_path = os.path.join(self._cache_path, key, mask.path) + + return data + def list_data(self) -> PagingList[RemoteData]: """List required Data object in a dataset segment.