Feature/Improved CLI lists with enhanced tables and new formats

pyproject.toml

@@ -1,7 +1,7 @@
 [tool.poetry]
 name = "zenml"
 version = "0.85.0"
-packages = [{ include = "zenml", from = "src" }]
+packages = [{ include = "zenml", from = "src" }, { include = "zenml_cli", from = "src" }]
 description = "ZenML: Write production-ready ML code."
 authors = ["ZenML GmbH <[email protected]>"]
 readme = "README.md"
@@ -38,7 +38,7 @@ exclude = [
 include = ["src/zenml", "*.txt", "*.sh", "*.md"]
 
 [tool.poetry.scripts]
-zenml = "zenml.cli.cli:cli"
+zenml = "zenml_cli:cli"
 
 [tool.poetry.dependencies]
 alembic = { version = ">=1.8.1,<=1.15.2" }

src/zenml/cli/artifact.py

@@ -13,18 +13,23 @@
 #  permissions and limitations under the License.
 """CLI functionality to interact with artifacts."""
 
-from typing import Any, Dict, List, Optional
+from typing import Any, List, Optional
 
 import click
 
 from zenml.cli import utils as cli_utils
 from zenml.cli.cli import TagGroup, cli
+from zenml.cli.utils import (
+    list_options,
+)
 from zenml.client import Client
+from zenml.console import console
 from zenml.enums import CliCategories
 from zenml.logger import get_logger
-from zenml.models import ArtifactFilter, ArtifactVersionFilter
-from zenml.models.v2.core.artifact import ArtifactResponse
-from zenml.models.v2.core.artifact_version import ArtifactVersionResponse
+from zenml.models import (
+    ArtifactFilter,
+    ArtifactVersionFilter,
+)
 from zenml.utils.pagination_utils import depaginate
 
 logger = get_logger(__name__)
@@ -35,25 +40,39 @@ def artifact() -> None:
     """Commands for interacting with artifacts."""
 
 
-@cli_utils.list_options(ArtifactFilter)
+@list_options(
+    ArtifactFilter,
+    default_columns=["id", "name", "latest_version_name", "user", "created"],
+)
 @artifact.command("list", help="List all artifacts.")
-def list_artifacts(**kwargs: Any) -> None:
+def list_artifacts(output_format: str, columns: str, **kwargs: Any) -> None:
     """List all artifacts.
 
     Args:
-        **kwargs: Keyword arguments to filter artifacts by.
+        output_format: Output format (table, json, yaml, tsv, csv).
+        columns: Comma-separated list of columns to display.
+        kwargs: Keyword arguments to filter artifacts by.
     """
-    artifacts = Client().list_artifacts(**kwargs)
-
-    if not artifacts:
-        cli_utils.declare("No artifacts found.")
-        return
-
-    to_print = []
-    for artifact in artifacts:
-        to_print.append(_artifact_to_print(artifact))
+    with console.status("Listing artifacts..."):
+        artifacts = Client().list_artifacts(**kwargs)
+
+        artifact_list = []
+        for artifact in artifacts.items:
+            artifact_data = cli_utils.prepare_response_data(artifact)
+            artifact_data.update(
+                {
+                    "latest_version_name": artifact.latest_version_name,
+                    "latest_version_id": artifact.latest_version_id,
+                }
+            )
+            artifact_list.append(artifact_data)
 
-    cli_utils.print_table(to_print)
+    cli_utils.handle_output(
+        artifact_list,
+        pagination_info=artifacts.pagination_info,
+        columns=columns,
+        output_format=output_format,
+    )
 
 
 @artifact.command("update", help="Update an artifact.")
@@ -115,25 +134,42 @@ def version() -> None:
     """Commands for interacting with artifact versions."""
 
 
-@cli_utils.list_options(ArtifactVersionFilter)
+@list_options(
+    ArtifactVersionFilter,
+    default_columns=["id", "name", "version", "type", "user", "created"],
+)
 @version.command("list", help="List all artifact versions.")
-def list_artifact_versions(**kwargs: Any) -> None:
+def list_artifact_versions(
+    output_format: str, columns: str, **kwargs: Any
+) -> None:
     """List all artifact versions.
 
     Args:
-        **kwargs: Keyword arguments to filter artifact versions by.
+        output_format: Output format (table, json, yaml, tsv, csv).
+        columns: Comma-separated list of columns to display.
+        kwargs: Keyword arguments to filter artifact versions by.
     """
-    artifact_versions = Client().list_artifact_versions(**kwargs)
-
-    if not artifact_versions:
-        cli_utils.declare("No artifact versions found.")
-        return
+    with console.status("Listing artifact versions..."):
+        artifact_versions = Client().list_artifact_versions(**kwargs)
 
-    to_print = []
-    for artifact_version in artifact_versions:
-        to_print.append(_artifact_version_to_print(artifact_version))
+        artifact_version_list = []
+        for artifact_version in artifact_versions.items:
+            artifact_version_data = cli_utils.prepare_response_data(
+                artifact_version
+            )
+            artifact_version_data.update(
+                {
+                    "name": artifact_version.artifact.name,
+                }
+            )
+            artifact_version_list.append(artifact_version_data)
 
-    cli_utils.print_table(to_print)
+    cli_utils.handle_output(
+        artifact_version_list,
+        pagination_info=artifact_versions.pagination_info,
+        columns=columns,
+        output_format=output_format,
+    )
 
 
 @version.command("update", help="Update an artifact version.")
@@ -294,28 +330,3 @@ def prune_artifacts(
                     f"Failed to delete artifact version {unused_artifact_version.id}: {str(e)}"
                 )
     cli_utils.declare("All unused artifacts and artifact versions deleted.")
-
-
-def _artifact_version_to_print(
-    artifact_version: ArtifactVersionResponse,
-) -> Dict[str, Any]:
-    return {
-        "id": artifact_version.id,
-        "name": artifact_version.artifact.name,
-        "version": artifact_version.version,
-        "uri": artifact_version.uri,
-        "type": artifact_version.type,
-        "materializer": artifact_version.materializer,
-        "data_type": artifact_version.data_type,
-        "tags": [t.name for t in artifact_version.tags],
-    }
-
-
-def _artifact_to_print(
-    artifact_version: ArtifactResponse,
-) -> Dict[str, Any]:
-    return {
-        "id": artifact_version.id,
-        "name": artifact_version.name,
-        "tags": [t.name for t in artifact_version.tags],
-    }

src/zenml/cli/authorized_device.py

@@ -56,27 +56,37 @@ def describe_authorized_device(id_or_prefix: str) -> None:
     )
 
 
+@list_options(
+    OAuthDeviceFilter,
+    default_columns=["status", "ip_address", "hostname", "os", "created"],
+)
 @authorized_device.command(
     "list", help="List all authorized devices for the current user."
 )
-@list_options(OAuthDeviceFilter)
-def list_authorized_devices(**kwargs: Any) -> None:
+def list_authorized_devices(
+    output_format: str, columns: str, **kwargs: Any
+) -> None:
     """List all authorized devices.
 
     Args:
-        **kwargs: Keyword arguments to filter authorized devices.
+        output_format: Output format (table, json, yaml, tsv, csv).
+        columns: Comma-separated list of columns to display.
+        kwargs: Keyword arguments to filter authorized devices.
     """
-    with console.status("Listing authorized devices...\n"):
+    with console.status("Listing authorized devices..."):
         devices = Client().list_authorized_devices(**kwargs)
 
-        if not devices.items:
-            cli_utils.declare("No authorized devices found for this filter.")
-            return
+        device_list = []
+        for device in devices.items:
+            device_data = cli_utils.prepare_response_data(device)
+            device_list.append(device_data)
 
-        cli_utils.print_pydantic_models(
-            devices,
-            columns=["id", "status", "ip_address", "hostname", "os"],
-        )
+    cli_utils.handle_output(
+        device_list,
+        pagination_info=devices.pagination_info,
+        columns=columns,
+        output_format=output_format,
+    )
 
 
 @authorized_device.command("lock")

src/zenml/cli/code_repository.py

@@ -19,7 +19,9 @@
 
 from zenml.cli import utils as cli_utils
 from zenml.cli.cli import TagGroup, cli
-from zenml.cli.utils import list_options
+from zenml.cli.utils import (
+    list_options,
+)
 from zenml.client import Client
 from zenml.code_repositories import BaseCodeRepository
 from zenml.config.source import Source
@@ -188,25 +190,34 @@ def describe_code_repository(name_id_or_prefix: str) -> None:
         )
 
 
+@list_options(
+    CodeRepositoryFilter, default_columns=["name", "type", "url", "created"]
+)
 @code_repository.command("list", help="List all connected code repositories.")
-@list_options(CodeRepositoryFilter)
-def list_code_repositories(**kwargs: Any) -> None:
+def list_code_repositories(
+    output_format: str, columns: str, **kwargs: Any
+) -> None:
     """List all connected code repositories.
 
     Args:
-        **kwargs: Keyword arguments to filter code repositories.
+        output_format: Output format (table, json, yaml, tsv, csv).
+        columns: Comma-separated list of columns to display.
+        kwargs: Keyword arguments to filter code repositories.
     """
-    with console.status("Listing code repositories...\n"):
+    with console.status("Listing code repositories..."):
         repos = Client().list_code_repositories(**kwargs)
 
-        if not repos.items:
-            cli_utils.declare("No code repositories found for this filter.")
-            return
+        repo_list = []
+        for repo in repos.items:
+            repo_data = cli_utils.prepare_response_data(repo)
+            repo_list.append(repo_data)
 
-        cli_utils.print_pydantic_models(
-            repos,
-            exclude_columns=["created", "updated", "user", "project"],
-        )
+    cli_utils.handle_output(
+        repo_list,
+        pagination_info=repos.pagination_info,
+        columns=columns,
+        output_format=output_format,
+    )
 
 
 @code_repository.command(