Feature/Improved CLI lists with enhanced tables and new formats

pyproject.toml

@@ -1,7 +1,7 @@
 [tool.poetry]
 name = "zenml"
 version = "0.84.2"
-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

@@ -19,7 +19,13 @@
 
 from zenml.cli import utils as cli_utils
 from zenml.cli.cli import TagGroup, cli
+from zenml.cli.utils import (
+    enhanced_list_options,
+    format_date_for_table,
+    prepare_list_data,
+)
 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
@@ -35,25 +41,39 @@ def artifact() -> None:
     """Commands for interacting with artifacts."""
 
 
-@cli_utils.list_options(ArtifactFilter)
+@enhanced_list_options(ArtifactFilter)
 @artifact.command("list", help="List all artifacts.")
 def list_artifacts(**kwargs: Any) -> None:
     """List all artifacts.
 
     Args:
         **kwargs: Keyword arguments to filter artifacts by.
     """
-    artifacts = Client().list_artifacts(**kwargs)
+    # Extract table options from kwargs
+    table_kwargs = cli_utils.extract_table_options(kwargs)
+
+    with console.status("Listing artifacts..."):
+        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))
+    # Prepare data based on output format
+    output_format = (
+        table_kwargs.get("output") or cli_utils.get_default_output_format()
+    )
+    artifact_data = []
+
+    # Use centralized data preparation
+    artifact_data = prepare_list_data(
+        artifacts.items, output_format, _artifact_to_print
+    )
 
-    cli_utils.print_table(to_print)
+    # Handle table output with enhanced system and pagination
+    cli_utils.handle_table_output(
+        artifact_data, page=artifacts, **table_kwargs
+    )
 
 
 @artifact.command("update", help="Update an artifact.")
@@ -115,25 +135,39 @@ def version() -> None:
     """Commands for interacting with artifact versions."""
 
 
-@cli_utils.list_options(ArtifactVersionFilter)
+@enhanced_list_options(ArtifactVersionFilter)
 @version.command("list", help="List all artifact versions.")
 def list_artifact_versions(**kwargs: Any) -> None:
     """List all artifact versions.
 
     Args:
         **kwargs: Keyword arguments to filter artifact versions by.
     """
-    artifact_versions = Client().list_artifact_versions(**kwargs)
+    # Extract table options from kwargs
+    table_kwargs = cli_utils.extract_table_options(kwargs)
+
+    with console.status("Listing artifact versions..."):
+        artifact_versions = Client().list_artifact_versions(**kwargs)
 
     if not artifact_versions:
         cli_utils.declare("No artifact versions found.")
         return
 
-    to_print = []
-    for artifact_version in artifact_versions:
-        to_print.append(_artifact_version_to_print(artifact_version))
+    # Prepare data based on output format
+    output_format = (
+        table_kwargs.get("output") or cli_utils.get_default_output_format()
+    )
+    artifact_version_data = []
 
-    cli_utils.print_table(to_print)
+    # Use centralized data preparation
+    artifact_version_data = prepare_list_data(
+        artifact_versions.items, output_format, _artifact_version_to_print
+    )
+
+    # Handle table output with enhanced system and pagination
+    cli_utils.handle_table_output(
+        artifact_version_data, page=artifact_versions, **table_kwargs
+    )
 
 
 @version.command("update", help="Update an artifact version.")
@@ -296,26 +330,71 @@ def prune_artifacts(
     cli_utils.declare("All unused artifacts and artifact versions deleted.")
 
 
+def _artifact_to_print(artifact: ArtifactResponse) -> Dict[str, Any]:
+    """Convert an artifact response to a dictionary suitable for table display.
+
+    For table output, keep it compact with essential artifact information.
+    Full details are available in JSON/YAML output formats.
+
+    Args:
+        artifact: Artifact response object
+
+    Returns:
+        Dictionary containing formatted artifact data for table display
+    """
+    return {
+        "name": artifact.name,
+        "tags": [t.name for t in artifact.tags] if artifact.tags else [],
+        "created": format_date_for_table(artifact.created),
+    }
+
+
+def _artifact_to_print_full(artifact: ArtifactResponse) -> Dict[str, Any]:
+    """Convert artifact response to complete dictionary for JSON/YAML.
+
+    Args:
+        artifact: Artifact response object
+
+    Returns:
+        Complete dictionary containing all artifact data
+    """
+    return artifact.model_dump(mode="json")
+
+
 def _artifact_version_to_print(
     artifact_version: ArtifactVersionResponse,
 ) -> Dict[str, Any]:
+    """Convert artifact version response to dictionary for table display.
+
+    For table output, keep it compact with essential version information.
+    Full details are available in JSON/YAML output formats.
+
+    Args:
+        artifact_version: Artifact version response object
+
+    Returns:
+        Dictionary containing formatted artifact version data for table display
+    """
     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],
+        "tags": [t.name for t in artifact_version.tags]
+        if artifact_version.tags
+        else [],
+        "created": format_date_for_table(artifact_version.created),
     }
 
 
-def _artifact_to_print(
-    artifact_version: ArtifactResponse,
+def _artifact_version_to_print_full(
+    artifact_version: ArtifactVersionResponse,
 ) -> Dict[str, Any]:
-    return {
-        "id": artifact_version.id,
-        "name": artifact_version.name,
-        "tags": [t.name for t in artifact_version.tags],
-    }
+    """Convert artifact version response to complete dictionary for JSON/YAML.
+
+    Args:
+        artifact_version: Artifact version response object
+
+    Returns:
+        Complete dictionary containing all artifact version data
+    """
+    return artifact_version.model_dump(mode="json")

src/zenml/cli/authorized_device.py

@@ -13,13 +13,17 @@
 #  permissions and limitations under the License.
 """CLI functionality to interact with authorized devices."""
 
-from typing import Any
+from typing import Any, Dict
 
 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.cli.utils import (
+    format_date_for_table,
+    list_options,
+    prepare_list_data,
+)
 from zenml.client import Client
 from zenml.console import console
 from zenml.enums import CliCategories
@@ -29,6 +33,26 @@
 logger = get_logger(__name__)
 
 
+def _authorized_device_to_print(device: Any) -> Dict[str, Any]:
+    """Convert an authorized device response to a dictionary for table display.
+
+    Args:
+        device: Authorized device response object
+
+    Returns:
+        Dictionary containing formatted device data for table display
+    """
+    return {
+        "status": device.status.value
+        if hasattr(device.status, "value")
+        else str(device.status),
+        "ip_address": device.ip_address or "",
+        "hostname": device.hostname or "",
+        "os": device.os or "",
+        "created": format_date_for_table(device.created),
+    }
+
+
 @cli.group(cls=TagGroup, tag=CliCategories.MANAGEMENT_TOOLS)
 def authorized_device() -> None:
     """Interact with authorized devices."""
@@ -66,16 +90,31 @@ def list_authorized_devices(**kwargs: Any) -> None:
     Args:
         **kwargs: Keyword arguments to filter authorized devices.
     """
-    with console.status("Listing authorized devices...\n"):
+    # Extract table options from kwargs
+    table_kwargs = cli_utils.extract_table_options(kwargs)
+
+    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.")
+            cli_utils.declare("No authorized devices found.")
             return
 
-        cli_utils.print_pydantic_models(
-            devices,
-            columns=["id", "status", "ip_address", "hostname", "os"],
+        # Prepare data based on output format
+        output_format = (
+            table_kwargs.get("output") or cli_utils.get_default_output_format()
+        )
+
+        # Use centralized data preparation
+        device_data = prepare_list_data(
+            devices.items, output_format, _authorized_device_to_print
+        )
+
+        # Handle table output with enhanced system and pagination
+        cli_utils.handle_table_output(
+            data=device_data,
+            page=devices,
+            **table_kwargs,
         )
 
 

src/zenml/cli/code_repository.py

@@ -13,13 +13,17 @@
 #  permissions and limitations under the License.
 """CLI functionality to interact with code repositories."""
 
-from typing import Any, List, Optional
+from typing import Any, Dict, 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.cli.utils import (
+    enhanced_list_options,
+    format_date_for_table,
+    prepare_list_data,
+)
 from zenml.client import Client
 from zenml.code_repositories import BaseCodeRepository
 from zenml.config.source import Source
@@ -32,6 +36,29 @@
 logger = get_logger(__name__)
 
 
+def _code_repository_to_print(repo: Any) -> Dict[str, Any]:
+    """Convert a code repository response to a dictionary for table display.
+
+    Args:
+        repo: Code repository response object
+
+    Returns:
+        Dictionary containing formatted repository data for table display
+    """
+    return {
+        "name": repo.name,
+        "type": repo.source.type
+        if hasattr(repo, "source") and repo.source
+        else "Unknown",
+        "url": repo.source.url
+        if hasattr(repo, "source")
+        and repo.source
+        and hasattr(repo.source, "url")
+        else "",
+        "created": format_date_for_table(repo.created),
+    }
+
+
 @cli.group(cls=TagGroup, tag=CliCategories.MANAGEMENT_TOOLS)
 def code_repository() -> None:
     """Interact with code repositories."""
@@ -189,23 +216,39 @@ def describe_code_repository(name_id_or_prefix: str) -> None:
 
 
 @code_repository.command("list", help="List all connected code repositories.")
-@list_options(CodeRepositoryFilter)
+@enhanced_list_options(CodeRepositoryFilter)
 def list_code_repositories(**kwargs: Any) -> None:
     """List all connected code repositories.
 
     Args:
         **kwargs: Keyword arguments to filter code repositories.
     """
-    with console.status("Listing code repositories...\n"):
+    # Extract table options from kwargs
+    table_kwargs = cli_utils.extract_table_options(kwargs)
+
+    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.")
+            cli_utils.declare("No code repositories found.")
             return
 
-        cli_utils.print_pydantic_models(
-            repos,
-            exclude_columns=["created", "updated", "user", "project"],
+        # Prepare data based on output format
+        output_format = (
+            table_kwargs.get("output") or cli_utils.get_default_output_format()
+        )
+        repo_data = []
+
+        # Use centralized data preparation
+        repo_data = prepare_list_data(
+            repos.items, output_format, _code_repository_to_print
+        )
+
+        # Handle table output with enhanced system
+        cli_utils.handle_table_output(
+            data=repo_data,
+            page=repos,
+            **table_kwargs,
         )