def _generate_segments(self,
offset: int = 0,
limit: int = 128) -> Generator[Segment, None, int]:
response = self._list_segments(offset, limit)
for item in response["segments"]:
segment = Segment._from_client( # pylint: disable=protected-access
SegmentClient(item["name"], self))
segment.description = item["description"]
yield segment
return response["totalCount"] # type: ignore[no-any-return]
def get_or_create_segment(self, name: str = "default") -> SegmentClient:
"""Get or create a segment with the given name.
Arguments:
name: The name of the fusion segment.
Returns:
The created :class:`~tensorbay.client.segment.SegmentClient` with given name.
"""
self._status.check_authority_for_draft()
if name not in self.list_segment_names():
self._create_segment(name)
return SegmentClient(name, self)
def test_upload_segment(self, mocker):
self.dataset_client._status.checkout(draft_number=1)
segment_test = Segment(name="test1")
for i in range(5):
segment_test.append(Data(f"data{i}.png"))
segment_client = SegmentClient(name="test1", data_client=self.dataset_client)
upload_segment = mocker.patch(
f"{dataset.__name__}.DatasetClient._upload_segment", return_value=segment_client
)
assert self.dataset_client.upload_segment(segment_test).name == "test1"
args, keywords = upload_segment.call_args
assert args[0] == segment_test
assert keywords["jobs"] == 1
assert not keywords["skip_uploaded_files"]
upload_segment.assert_called_once()
def test__upload_segment(self, mocker):
segment_test = Segment(name="test1")
for i in range(5):
segment_test.append(Data(f"data{i}.png"))
segment_client = SegmentClient(name="test1",
data_client=self.dataset_client)
get_or_create_segment = mocker.patch(
f"{dataset.__name__}.DatasetClient.get_or_create_segment",
return_value=segment_client)
list_data_paths = mocker.patch(
f"{segment.__name__}.SegmentClient.list_data_paths",
return_value=["data1.png", "data2.png"],
)
multithread_upload = mocker.patch(
f"{dataset.__name__}.multithread_upload")
with Tqdm(5, disable=False) as pbar:
self.dataset_client._upload_segment(segment_test,
skip_uploaded_files=True,
pbar=pbar)
get_or_create_segment.assert_called_once_with(segment_test.name)
list_data_paths.assert_called_once_with()
args, keywords = multithread_upload.call_args
assert args[0] == segment_client._upload_or_import_data
assert [item.path for item in args[1]
] == ["data0.png", "data3.png", "data4.png"]
assert keywords[
"callback"] == segment_client._synchronize_upload_info
assert keywords["jobs"] == 1
assert keywords["pbar"] == pbar
multithread_upload.assert_called_once()
with Tqdm(5, disable=False) as pbar:
self.dataset_client._upload_segment(segment_test,
skip_uploaded_files=False,
pbar=pbar)
get_or_create_segment.assert_called_with(segment_test.name)
list_data_paths.assert_called_with()
args, keywords = multithread_upload.call_args
assert args[0] == segment_client._upload_or_import_data
assert [item.path
for item in args[1]] == [f"data{i}.png" for i in range(5)]
assert keywords[
"callback"] == segment_client._synchronize_upload_info
assert keywords["jobs"] == 1
assert keywords["pbar"] == pbar
multithread_upload.assert_called()
def get_segment(self, name: str = "default") -> SegmentClient:
"""Get a segment in a certain commit according to given name.
Arguments:
name: The name of the required segment.
Returns:
The required :class:`~tensorbay.client.segment.SegmentClient`.
Raises:
ResourceNotExistError: When the required segment does not exist.
"""
if name not in self.list_segment_names():
raise ResourceNotExistError(resource="segment",
identification=name)
return SegmentClient(name, self)
def create_segment(self, name: str = "default") -> SegmentClient:
"""Create a segment with the given name.
Arguments:
name: The name of the fusion segment.
Returns:
The created :class:`~tensorbay.client.segment.SegmentClient` with given name.
Raises:
NameConflictError: When the segment exists.
"""
self._status.check_authority_for_draft()
if name not in self.list_segment_names():
self._create_segment(name)
else:
raise NameConflictError(resource="segment", identification=name)
return SegmentClient(name, self)
def copy_segment(
self,
source_name: str,
target_name: Optional[str] = None,
*,
source_client: Optional["DatasetClient"] = None,
strategy: str = "abort",
) -> SegmentClient:
"""Copy segment to this dataset.
Arguments:
source_name: The source name of the copied segment.
target_name: The target name of the copied segment.
This argument is used to specify a new name of the copied segment.
If None, the name of the copied segment will not be changed after copy.
source_client: The source dataset client of the copied segment.
This argument is used to specify where the copied segment comes from when the
copied segment is from another commit, draft or even another dataset.
If None, the copied segment comes from this dataset.
strategy: The strategy of handling the name conflict. There are three options:
1. "abort": stop copying and raise exception;
2. "override": the source segment will override the origin segment;
3. "skip": keep the origin segment.
Returns:
The client of the copied target segment.
"""
if not target_name:
target_name = source_name
self._copy_segment(source_name,
target_name,
source_client=source_client,
strategy=strategy)
return SegmentClient(target_name, self)
def move_segment(
self,
source_name: str,
target_name: str,
*,
strategy: str = "abort",
) -> SegmentClient:
"""Move/Rename segment in this dataset.
Arguments:
source_name: The source name of the moved segment.
target_name: The target name of the moved segment.
strategy: The strategy of handling the name conflict. There are three options:
1. "abort": stop moving and raise exception;
2. "override": the source segment will override the origin segment;
3. "skip": keep the origin segment.
Returns:
The client of the moved target segment.
"""
self._move_segment(source_name, target_name, strategy=strategy)
return SegmentClient(target_name, self)