feat(sdk): add autocomplete and version options to kfp cli (#7567)

* add helpful options to cli

* add tests

sdk/python/kfp/cli/cli.py

@@ -12,9 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import os
 from itertools import chain
 
 import click
+import kfp
 from kfp.cli import component
 from kfp.cli import diagnose_me_cli
 from kfp.cli import experiment
@@ -33,10 +35,39 @@ COMMANDS = {
     'no_client': {diagnose_me_cli.diagnose_me, component.component}
 }
 
+PROGRAM_NAME = 'kfp'
+
+SHELL_FILES = {
+    'bash': ['.bashrc'],
+    'zsh': ['.zshrc'],
+    'fish': ['.config', 'fish', 'completions', f'{PROGRAM_NAME}.fish']
+}
+
+
+def _create_completion(shell: str) -> str:
+    return f'eval "$(_{PROGRAM_NAME.upper()}_COMPLETE={shell}_source {PROGRAM_NAME})"'
+
+
+def _install_completion(shell: str) -> None:
+    completion_statement = _create_completion(shell)
+    source_file = os.path.join(os.path.expanduser('~'), *SHELL_FILES[shell])
+    with open(source_file, 'a') as f:
+        f.write('\n' + completion_statement + '\n')
+
 
 @click.group(
-    cls=aliased_plurals_group.AliasedPluralsGroup,
-    commands=list(chain.from_iterable(COMMANDS.values())))  # type: ignore
+    name=PROGRAM_NAME,
+    cls=aliased_plurals_group.AliasedPluralsGroup,  # type: ignore
+    commands=list(chain.from_iterable(COMMANDS.values())),  # type: ignore
+    invoke_without_command=True)
+@click.option(
+    '--show-completion',
+    type=click.Choice(list(SHELL_FILES.keys())),
+    default=None)
+@click.option(
+    '--install-completion',
+    type=click.Choice(list(SHELL_FILES.keys())),
+    default=None)
 @click.option('--endpoint', help='Endpoint of the KFP API service to connect.')
 @click.option('--iap-client-id', help='Client ID for IAP protected endpoint.')
 @click.option(
@@ -58,21 +89,31 @@ COMMANDS = {
     show_default=True,
     help='The formatting style for command output.')
 @click.pass_context
+@click.version_option(version=kfp.__version__, message='%(prog)s %(version)s')
 def cli(ctx: click.Context, endpoint: str, iap_client_id: str, namespace: str,
-        other_client_id: str, other_client_secret: str, output: OutputFormat):
+        other_client_id: str, other_client_secret: str, output: OutputFormat,
+        show_completion: str, install_completion: str):
     """kfp is the command line interface to KFP service.
 
     Feature stage:
     [Alpha](https://github.com/kubeflow/pipelines/blob/07328e5094ac2981d3059314cc848fbb71437a76/docs/release/feature-stages.md#alpha)
     """
+    if show_completion:
+        click.echo(_create_completion(show_completion))
+        return
+    if install_completion:
+        _install_completion(install_completion)
+        return
+
     client_commands = set(
         chain.from_iterable([
             (command.name, f'{command.name}s')
             for command in COMMANDS['client']  # type: ignore
         ]))
-
-    if ctx.invoked_subcommand in client_commands:
-        ctx.obj['client'] = Client(endpoint, iap_client_id, namespace,
-                                   other_client_id, other_client_secret)
-        ctx.obj['namespace'] = namespace
-        ctx.obj['output'] = output
+    if ctx.invoked_subcommand not in client_commands:
+        # Do not create a client for these subcommands
+        return
+    ctx.obj['client'] = Client(endpoint, iap_client_id, namespace,
+                               other_client_id, other_client_secret)
+    ctx.obj['namespace'] = namespace
+    ctx.obj['output'] = output

sdk/python/kfp/cli/cli_test.py

@@ -13,13 +13,19 @@
 # limitations under the License.
 
 import functools
+import itertools
+import os
+import re
+import tempfile
 import unittest
+from unittest import mock
 
+from absl.testing import parameterized
 from click import testing
 from kfp.cli import cli
 
 
-class TestCli(unittest.TestCase):
+class TestCliNounAliases(unittest.TestCase):
 
     def setUp(self):
         runner = testing.CliRunner()
@@ -39,3 +45,71 @@ class TestCli(unittest.TestCase):
         self.assertEqual(result.exit_code, 2)
         self.assertEqual("Error: Unrecognized command 'componentss'\n",
                          result.output)
+
+
+class TestCliAutocomplete(parameterized.TestCase):
+
+    def setUp(self):
+        runner = testing.CliRunner()
+        self.invoke = functools.partial(
+            runner.invoke, cli=cli.cli, catch_exceptions=False, obj={})
+
+    @parameterized.parameters(['bash', 'zsh', 'fish'])
+    def test_show_autocomplete(self, shell):
+        result = self.invoke(args=['--show-completion', shell])
+        expected = cli._create_completion(shell)
+        self.assertTrue(expected in result.output)
+        self.assertEqual(result.exit_code, 0)
+
+    @parameterized.parameters(['bash', 'zsh', 'fish'])
+    def test_install_autocomplete_with_empty_file(self, shell):
+        with tempfile.TemporaryDirectory() as tempdir:
+            with mock.patch('os.path.expanduser', return_value=tempdir):
+                temp_path = os.path.join(tempdir, *cli.SHELL_FILES[shell])
+                os.makedirs(os.path.dirname(temp_path), exist_ok=True)
+
+                result = self.invoke(args=['--install-completion', shell])
+                expected = cli._create_completion(shell)
+
+                with open(temp_path) as f:
+                    last_line = f.readlines()[-1]
+                self.assertEqual(expected + '\n', last_line)
+                self.assertEqual(result.exit_code, 0)
+
+    @parameterized.parameters(
+        list(itertools.product(['bash', 'zsh', 'fish'], [True, False])))
+    def test_install_autocomplete_with_unempty_file(self, shell,
+                                                    has_trailing_newline):
+        with tempfile.TemporaryDirectory() as tempdir:
+            with mock.patch('os.path.expanduser', return_value=tempdir):
+                temp_path = os.path.join(tempdir, *cli.SHELL_FILES[shell])
+                os.makedirs(os.path.dirname(temp_path), exist_ok=True)
+
+                existing_file_contents = [
+                    "something\n",
+                    "something else" + ('\n' if has_trailing_newline else ''),
+                ]
+                with open(temp_path, 'w') as f:
+                    f.writelines(existing_file_contents)
+
+                result = self.invoke(args=['--install-completion', shell])
+                expected = cli._create_completion(shell)
+
+                with open(temp_path) as f:
+                    last_line = f.readlines()[-1]
+                self.assertEqual(expected + '\n', last_line)
+                self.assertEqual(result.exit_code, 0)
+
+
+class TestCliVersion(unittest.TestCase):
+
+    def setUp(self):
+        runner = testing.CliRunner()
+        self.invoke = functools.partial(
+            runner.invoke, cli=cli.cli, catch_exceptions=False, obj={})
+
+    def test_version(self):
+        result = self.invoke(args=['--version'])
+        self.assertEqual(result.exit_code, 0)
+        matches = re.match(r'^kfp \d\.\d\.\d.*', result.output)
+        self.assertTrue(matches)

sdk/python/kfp/components/structures.py

@@ -15,7 +15,6 @@
 
 import dataclasses
 import itertools
-import json
 from typing import Any, Dict, Mapping, Optional, Sequence, Union
 
 import pydantic
@@ -596,4 +595,4 @@ class ComponentSpec(BaseModel):
         Args:
             output_file: File path to store the component yaml.
         """
-        ir_utils._write_ir_to_file(self.dict(), output_file)
+        ir_utils._write_ir_to_file(self.dict(), output_file)