Merge 3d48155 into a4f2b66

Author: yuyouyu32

tensorbay/cli/cli.py

@@ -168,6 +168,7 @@ def dataset(obj: Dict[str, str], tbrn: str, is_delete: bool, yes: bool) -> None:
 )
 @click.argument("tbrn", type=str)
 @click.option("-l", "--list", "is_list", is_flag=True, help="List the drafts.")
+@click.option("-e", "--edit", is_flag=True, help="Edit the draft's title and description.")
 @click.option(
     "-m",
     "--message",
@@ -188,6 +189,7 @@ def draft(
     obj: Dict[str, str],
     tbrn: str,
     is_list: bool,
+    edit: bool,
     message: Tuple[str, ...],
 ) -> None:
     """List or create drafts.\f
@@ -196,12 +198,13 @@ def draft(
         obj: A dict contains config information.
         tbrn: The tbrn of the dataset.
         is_list: Whether to list the drafts.
+        edit: Whether to edit the draft's title and description.
         message: The message of the draft.
 
     """  # noqa: D301,D415
     from .draft import _implement_draft
 
-    _implement_draft(obj, tbrn, is_list, message)
+    _implement_draft(obj, tbrn, is_list, edit, message)
 
 
 @command(

tensorbay/cli/draft.py

@@ -15,15 +15,16 @@
 from .tbrn import TBRN, TBRNType
 from .utility import edit_input, error, get_dataset_client, get_gas
 
-_DRAFT_HINT = """
+_DRAFT_HINT = """{}
+{}
 # Please enter the title for your draft.
 # Lines starting with '#' will be ignored.
 # And an empty draft title aborts the creation.
 """
 
 
 def _implement_draft(
-    obj: Dict[str, str], tbrn: str, is_list: bool, message: Tuple[str, ...]
+    obj: Dict[str, str], tbrn: str, is_list: bool, edit: bool, message: Tuple[str, ...]
 ) -> None:
     gas = get_gas(**obj)
     info = TBRN(tbrn=tbrn)
@@ -34,6 +35,8 @@ def _implement_draft(
 
     if is_list:
         _list_drafts(dataset_client, info)
+    elif edit:
+        _edit_draft(dataset_client, info, message)
     else:
         _create_draft(dataset_client, info, message)
 
@@ -83,7 +86,7 @@ def _echo_draft(
     try:
         branch = dataset_client.get_branch(branch_name)
     except ResourceNotExistError:
-        error('The branch "{branch_name}" does not exist')
+        error(f'The branch "{branch_name}" does not exist')
 
     if branch.commit_id != ROOT_COMMIT_ID:
         click.echo(f"Branch: {branch_name}({branch.commit_id})")
@@ -95,3 +98,27 @@ def _echo_draft(
     click.echo(f"\n    {title}\n")
     if description:
         click.echo(f"    {description}\n")
+
+
+def _edit_draft(dataset_client: DatasetClientType, info: TBRN, message: Tuple[str, ...]) -> None:
+    if not info.is_draft:
+        error("Draft number is needed when editing draft")
+
+    if message:
+        title, description = message[0], "\n".join(message[1:])
+    else:
+        draft = dataset_client.get_draft(draft_number=info.draft_number)
+        title, description = edit_input(
+            _DRAFT_HINT.format(
+                draft.title, "\n" + draft.description + "\n" if draft.description else ""
+            )
+        )
+
+    if not title:
+        error("Aborting creating draft due to empty draft title")
+
+    dataset_client.update_draft(
+        draft_number=info.draft_number, title=title, description=description
+    )
+    click.echo(f"{info.get_tbrn()} update successful!")
+    _echo_draft(dataset_client, title, description, dataset_client.status.branch_name)

tensorbay/cli/utility.py

@@ -152,7 +152,7 @@ def edit_input(hint: str) -> Tuple[str, str]:
     config_parser = read_config()
 
     editor = config_parser["config"].get("editor") if config_parser.has_section("config") else None
-    input_info = click.edit(hint, editor=editor)
+    input_info = click.edit(hint, editor=editor, require_save=False)
     return _clean_up(input_info)