Feature/Improved CLI lists with enhanced tables and new formats

pyproject.toml

@@ -1,7 +1,7 @@
 [tool.poetry]
 name = "zenml"
 version = "0.84.3"
-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,12 +19,19 @@
 
 from zenml.cli import utils as cli_utils
 from zenml.cli.cli import TagGroup, cli
+from zenml.cli.utils import (
+    enhanced_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,
+    ArtifactResponse,
+    ArtifactVersionFilter,
+    ArtifactVersionResponse,
+)
 from zenml.utils.pagination_utils import depaginate
 
 logger = get_logger(__name__)
@@ -35,25 +42,55 @@ def artifact() -> None:
     """Commands for interacting with artifacts."""
 
 
-@cli_utils.list_options(ArtifactFilter)
+@enhanced_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:
     """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))
+    def enrichment_func(
+        item: ArtifactResponse, result: Dict[str, Any]
+    ) -> Dict[str, Any]:
+        """Enrich the artifact data with the latest version name and id.
+
+        Args:
+            item: The artifact response.
+            result: The result dictionary.
+
+        Returns:
+            The enriched result dictionary.
+        """
+        result.update(
+            {
+                "latest_version_name": item.latest_version_name,
+                "latest_version_id": item.latest_version_id,
+            }
+        )
+        return result
 
-    cli_utils.print_table(to_print)
+    artifact_data = cli_utils.prepare_data_from_responses(
+        artifacts.items, enrichment_func=enrichment_func
+    )
+
+    # 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 +152,50 @@ def version() -> None:
     """Commands for interacting with artifact versions."""
 
 
-@cli_utils.list_options(ArtifactVersionFilter)
+@enhanced_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:
     """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))
+    def enrichment_func(
+        item: ArtifactVersionResponse, result: Dict[str, Any]
+    ) -> Dict[str, Any]:
+        """Enrich the artifact version data with the artifact name.
+
+        Args:
+            item: The artifact version response.
+            result: The result dictionary.
+
+        Returns:
+            The enriched result dictionary.
+        """
+        result.update({"name": item.artifact.name})
+        return result
+
+    artifact_version_data = cli_utils.prepare_data_from_responses(
+        artifact_versions.items, enrichment_func=enrichment_func
+    )
 
-    cli_utils.print_table(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.")
@@ -294,28 +356,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

@@ -19,7 +19,10 @@
 
 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,
+    prepare_data_from_responses,
+)
 from zenml.client import Client
 from zenml.console import console
 from zenml.enums import CliCategories
@@ -56,26 +59,37 @@ def describe_authorized_device(id_or_prefix: str) -> None:
     )
 
 
+@enhanced_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:
     """List all authorized devices.
 
     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"],
+        # Use centralized data preparation
+        device_data = prepare_data_from_responses(devices.items)
+
+        # 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

@@ -19,7 +19,10 @@
 
 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,
+    prepare_data_from_responses,
+)
 from zenml.client import Client
 from zenml.code_repositories import BaseCodeRepository
 from zenml.config.source import Source
@@ -188,24 +191,34 @@ def describe_code_repository(name_id_or_prefix: str) -> None:
         )
 
 
+@enhanced_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:
     """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"],
+        # Use centralized data preparation
+        repo_data = prepare_data_from_responses(repos.items)
+
+        # Handle table output with enhanced system
+        cli_utils.handle_table_output(
+            data=repo_data,
+            page=repos,
+            **table_kwargs,
         )
 
 

src/zenml/cli/config.py

@@ -11,7 +11,7 @@
 #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
 #  or implied. See the License for the specific language governing
 #  permissions and limitations under the License.
-"""CLI for manipulating ZenML local and global config file."""
+"""CLI for ZenML analytics and logging configuration."""
 
 import click