def _implement_dataset(obj: ContextInfo, tbrn: str, is_delete: bool, yes: bool) -> None:
gas = obj.get_gas()
if not tbrn:
if is_delete:
error("Missing argument TBRN")
for dataset_name in gas.list_dataset_names():
click.echo(TBRN(dataset_name).get_tbrn())
return
tbrn_info = TBRN(tbrn=tbrn)
if tbrn_info.type != TBRNType.DATASET:
error(f'"{tbrn}" is not a dataset')
colored_tbrn = tbrn_info.get_colored_tbrn()
if is_delete:
if not yes:
click.confirm(
f'Dataset "{colored_tbrn}" will be completely deleted.\nDo you want to continue?',
abort=True,
)
gas.delete_dataset(tbrn_info.dataset_name)
click.echo(f'Successfully deleted dataset "{colored_tbrn}"')
return
gas.create_dataset(tbrn_info.dataset_name)
click.echo(f'Successfully created dataset "{colored_tbrn}"')
def _implement_ls(obj: ContextInfo, tbrn: str, list_all_files: bool,
show_total_num: bool) -> None:
gas = obj.get_gas()
if not tbrn:
dataset_names = gas.list_dataset_names()
if show_total_num:
click.echo(f"total {len(dataset_names)}")
for dataset_name in dataset_names:
click.echo(TBRN(dataset_name).get_tbrn())
return
tbrn_info = TBRN(tbrn=tbrn)
_LS_FUNCS[tbrn_info.type](gas, tbrn_info, list_all_files, show_total_num)
def _ls_dataset(gas: GAS, tbrn_info: TBRN, list_all_files: bool,
show_total_num: bool) -> None:
dataset_client = get_dataset_client(gas, tbrn_info)
segment_names = dataset_client.list_segment_names()
if not list_all_files:
if show_total_num:
click.echo(f"total {len(segment_names)}")
for segment_name in segment_names:
click.echo(TBRN(tbrn_info.dataset_name, segment_name).get_tbrn())
return
if isinstance(dataset_client, FusionDatasetClient):
error('"-a" flag is not supported for fusion dataset yet')
segment_paths = [
dataset_client.get_segment(segment_name).list_data_paths()
for segment_name in segment_names
]
if show_total_num:
total = sum(len(segment_path) for segment_path in segment_paths)
click.echo(f"total {total}")
for segment_name, segment_path in zip(segment_names, segment_paths):
_echo_data(
tbrn_info.dataset_name,
tbrn_info.draft_number,
tbrn_info.revision,
segment_name,
segment_path,
)
def _echo_data(
dataset_name: str,
draft_number: Optional[int],
revision: Optional[str],
segment_name: str,
data_iter: Iterable[str],
) -> None:
"""Echo files in data_iter under 'tb:dataset_name:segment_name'.
Arguments:
dataset_name: The name of the dataset the segment belongs to.
draft_number: The draft number (if the status is draft).
revision: The revision (if the status is not draft).
segment_name: The name of the segment.
data_iter: Iterable data to be echoed.
"""
for data in data_iter:
click.echo(
TBRN(
dataset_name,
segment_name,
remote_path=data,
draft_number=draft_number,
revision=revision,
).get_tbrn())
def _create_branch(dataset_client: DatasetClientType, name: str) -> None:
if dataset_client.status.is_draft:
error("Branch cannot be created from a draft")
dataset_client.create_branch(name)
branch_tbrn = TBRN(dataset_client.name, revision=name).get_colored_tbrn()
click.echo(f'Successfully created branch "{branch_tbrn}"')
def _implement_log( # pylint: disable=too-many-arguments
obj: ContextInfo,
tbrn: str,
max_count: Optional[int],
oneline: bool,
is_all: bool,
graph: bool,
show_drafts: bool,
) -> None:
gas = obj.get_gas()
tbrn_info = TBRN(tbrn=tbrn)
if tbrn_info.type != TBRNType.DATASET:
error(f'To log commits, "{tbrn}" must be a dataset')
dataset_client = gas._get_dataset_with_any_type( # pylint: disable=protected-access
tbrn_info.dataset_name
)
commit_id_to_branches: DefaultDict[str, List[str]] = defaultdict(list)
for branch in dataset_client.list_branches():
commit_id_to_branches[branch.commit_id].append(branch.name)
if is_all:
revisions: List[Optional[str]] = [branch.name for branch in dataset_client.list_branches()]
else:
revisions = (
[tbrn_info.revision] if tbrn_info.revision else [dataset_client.status.branch_name]
)
Printer: Union[Type[_GraphPrinter], Type[_Printer]] = _GraphPrinter if graph else _Printer
message_generator = islice(
Printer(
dataset_client, revisions, commit_id_to_branches, oneline, show_drafts=show_drafts
).generate_commits_and_drafts_messages(),
max_count,
)
_echo_messages(message_generator)
def _delete_tag(dataset_client: DatasetClientType, tbrn_info: TBRN) -> None:
if not tbrn_info.revision:
error(f'To delete a tag, "{tbrn_info.get_tbrn()}" must have a tag name')
dataset_client.delete_tag(tbrn_info.revision)
tag_tbrn = TBRN(dataset_client.name, revision=tbrn_info.revision).get_colored_tbrn()
click.echo(f'Successfully deleted tag "{tag_tbrn}"')
def _implement_cp( # pylint: disable=too-many-arguments
obj: ContextInfo,
local_paths: Iterable[str],
tbrn: str,
is_recursive: bool,
jobs: int,
skip_uploaded_files: bool,
) -> None:
gas = obj.get_gas()
tbrn_info = TBRN(tbrn=tbrn)
dataset_client = get_dataset_client(gas, tbrn_info, is_fusion=False)
if tbrn_info.type not in (TBRNType.SEGMENT, TBRNType.NORMAL_FILE):
error(f'"{tbrn}" is not a segment or file type')
target_remote_path = tbrn_info.remote_path if tbrn_info.type == TBRNType.NORMAL_FILE else ""
local_abspaths = [os.path.abspath(local_path) for local_path in local_paths]
if (
len(local_abspaths) == 1
and not os.path.isdir(local_abspaths[0])
and target_remote_path
and not target_remote_path.endswith("/")
):
segment_client = dataset_client.get_or_create_segment(tbrn_info.segment_name)
segment_client.upload_file(local_abspaths[0], target_remote_path)
else:
segment = _get_segment(
tbrn_info.segment_name, local_abspaths, target_remote_path, is_recursive
)
dataset_client.upload_segment(
segment, jobs=jobs, skip_uploaded_files=skip_uploaded_files, _is_cli=True
)
def _create_tag(dataset_client: DatasetClientType, name: str) -> None:
if dataset_client.status.is_draft:
error(f'To create a tag, "{dataset_client.name}" cannot be in draft status')
if not dataset_client.status.commit_id:
error(f'To create a tag, "{dataset_client.name}" should have commit history')
dataset_client.create_tag(name=name)
tag_tbrn = TBRN(dataset_client.name, revision=name).get_colored_tbrn()
click.echo(f'Successfully created tag "{tag_tbrn}"')
def _implement_commit(obj: ContextInfo, tbrn: str, message: Tuple[str, ...]) -> None:
gas = obj.get_gas()
tbrn_info = TBRN(tbrn=tbrn)
dataset_client = get_dataset_client(gas, tbrn_info)
if tbrn_info.type != TBRNType.DATASET:
error(f'To operate a commit, "{tbrn}" must be a dataset')
if not tbrn_info.is_draft:
error(f'To commit, "{tbrn}" must be in draft status, like "{tbrn}#1"')
dataset_client.checkout(draft_number=tbrn_info.draft_number)
draft = dataset_client.get_draft()
hint_message = format_hint(draft.title, draft.description, _COMMIT_HINT)
title, description = edit_message(message, hint_message, obj.config_parser)
if not title:
error("Aborting commit due to empty commit message")
dataset_client.commit(title, description)
commit_tbrn = TBRN(tbrn_info.dataset_name, revision=dataset_client.status.commit_id)
click.echo(
"Committed successfully: "
f"{tbrn_info.get_colored_tbrn()} -> {commit_tbrn.get_colored_tbrn()}"
)
def _implement_tag(obj: ContextInfo, tbrn: str, name: str, is_delete: bool) -> None:
tbrn_info = TBRN(tbrn=tbrn)
if tbrn_info.type != TBRNType.DATASET:
error(f'To operate a tag, "{tbrn}" must be a dataset')
gas = obj.get_gas()
dataset_client = get_dataset_client(gas, tbrn_info)
if is_delete:
_delete_tag(dataset_client, tbrn_info)
elif name:
_create_tag(dataset_client, name)
else:
_list_tags(dataset_client)
def _implement_branch(obj: ContextInfo, tbrn: str, name: str, verbose: bool,
is_delete: bool) -> None:
tbrn_info = TBRN(tbrn=tbrn)
if tbrn_info.type != TBRNType.DATASET:
error(f'To operate a branch, "{tbrn_info}" must be a dataset')
gas = obj.get_gas()
dataset_client = get_dataset_client(gas, tbrn_info)
if is_delete:
_delete_branch(dataset_client, tbrn_info)
return
if name:
_create_branch(dataset_client, name)
else:
_list_branches(dataset_client, verbose)
def _create_draft(
dataset_client: DatasetClientType,
tbrn_info: TBRN,
message: Tuple[str, ...],
config_parser: ConfigParser,
) -> None:
if tbrn_info.is_draft:
error(
f'Create a draft in draft status "{tbrn_info.get_tbrn()}" is not permitted'
)
title, description = edit_message(message, _DRAFT_HINT, config_parser)
if not title:
error("Aborting creating draft due to empty draft message")
dataset_client.create_draft(title=title, description=description)
status = dataset_client.status
draft_tbrn = TBRN(tbrn_info.dataset_name,
draft_number=status.draft_number).get_colored_tbrn()
click.echo(f'Successfully created draft "{draft_tbrn}"')
_echo_draft(dataset_client, title, description, status.branch_name)
def _implement_rm(obj: ContextInfo, tbrn: str, is_recursive: bool) -> None:
gas = obj.get_gas()
tbrn_info = TBRN(tbrn=tbrn)
dataset_client = get_dataset_client(gas, tbrn_info, is_fusion=False)
if tbrn_info.type not in (TBRNType.SEGMENT, TBRNType.NORMAL_FILE):
error(f'"{tbrn}" is an invalid path to remove')
if not tbrn_info.is_draft:
error(
f'To remove the data, "{tbrn}" must be in draft status, like "{tbrn}#1"'
)
if tbrn_info.type == TBRNType.SEGMENT:
if not is_recursive:
error("Please use -r option to remove the whole segment")
dataset_client.delete_segment(tbrn_info.segment_name)
else:
segment = dataset_client.get_segment(tbrn_info.segment_name)
segment.delete_data(tbrn_info.remote_path)
click.echo(f'Successfully deleted "{tbrn_info.get_colored_tbrn()}"')
def _implement_draft( # pylint: disable=too-many-arguments
obj: ContextInfo,
tbrn: str,
is_list: bool,
edit: bool,
close: bool,
message: Tuple[str, ...],
) -> None:
gas = obj.get_gas()
tbrn_info = TBRN(tbrn=tbrn)
dataset_client = get_dataset_client(gas, tbrn_info)
if tbrn_info.type != TBRNType.DATASET:
error(f'To operate a draft, "{tbrn}" must be a dataset')
if is_list:
_list_drafts(dataset_client, tbrn_info)
elif edit:
_edit_draft(dataset_client, tbrn_info, message, obj.config_parser)
elif close:
_close_draft(dataset_client, tbrn_info)
else:
_create_draft(dataset_client, tbrn_info, message, obj.config_parser)
