From ae4f71fe0c8d0c114322419fad1d33196f2d8d21 Mon Sep 17 00:00:00 2001 From: Alex Strick van Linschoten Date: Fri, 10 Oct 2025 16:52:29 +0200 Subject: [PATCH 01/25] Try fix for click 8.2 issue --- src/zenml/cli/formatter.py | 174 ++++++++++++++++++------------------- 1 file changed, 87 insertions(+), 87 deletions(-) diff --git a/src/zenml/cli/formatter.py b/src/zenml/cli/formatter.py index 1d35c757d36..3d0410d57ff 100644 --- a/src/zenml/cli/formatter.py +++ b/src/zenml/cli/formatter.py @@ -13,7 +13,7 @@ # permissions and limitations under the License. """Helper functions to format output for CLI.""" -from typing import Dict, Iterable, Iterator, Optional, Sequence, Tuple +from typing import Dict, Iterable, Iterator, List, Optional, Sequence, Tuple from click import formatting from click._compat import term_len @@ -79,105 +79,105 @@ def write_dl( col_max: int = 30, col_spacing: int = 2, ) -> None: - """Writes a definition list into the buffer. + """Writes a definition list into the formatter buffer. - This is how options and commands are usually formatted. + Click 8.2 tightened validation so that definition list entries must be + pairs. Our CLI groups commands in tagged triples, so we detect that + case and render it manually while delegating the standard behaviour to + Click for the classic two-column output. This keeps the formatter + compatible with both older Click releases (<8.2) and the newer ones + without sacrificing the custom grouped layout. - Arguments: - rows: a list of items as tuples for the terms and values. - col_max: the maximum width of the first column. - col_spacing: the number of spaces between the first and - second column (and third). - - The default behavior is to format the rows in a definition list - with rows of 2 columns following the format ``(term, value)``. - But for new CLI commands, we want to format the rows in a definition - list with rows of 3 columns following the format - ``(term, value, description)``. + Args: + rows: A sequence of tuples that represent the definition list + entries. + col_max: The maximum width of the first column. + col_spacing: The number of spaces between columns. Raises: - TypeError: if the number of columns is not 2 or 3. + TypeError: If the provided rows do not represent two or three + column entries. """ - rows = list(rows) + + normalized_rows: List[Tuple[str, ...]] = [tuple(row) for row in rows] + + if not normalized_rows: + return + + unique_lengths = {len(row) for row in normalized_rows} + + if unique_lengths == {2}: + two_col_rows = [(row[0], row[1]) for row in normalized_rows] + super().write_dl( + two_col_rows, col_max=col_max, col_spacing=col_spacing + ) + return + + if unique_lengths == {3}: + self._write_triple_definition_list( + [(row[0], row[1], row[2]) for row in normalized_rows], + col_max=col_max, + col_spacing=col_spacing, + ) + return + + raise TypeError( + "Expected either two- or three-column definition list entries." + ) + + def _write_triple_definition_list( + self, + rows: List[Tuple[str, str, str]], + col_max: int, + col_spacing: int, + ) -> None: widths = measure_table(rows) - if len(widths) == 2: - first_col = min(widths[0], col_max) + col_spacing - - for first, second in iter_rows(rows, len(widths)): - self.write(f"{'':>{self.current_indent}}{first}") - if not second: - self.write("\n") - continue - if term_len(first) <= first_col - col_spacing: - self.write(" " * (first_col - term_len(first))) - else: - self.write("\n") - self.write(" " * (first_col + self.current_indent)) - - text_width = max(self.width - first_col - 2, 10) - wrapped_text = formatting.wrap_text( - second, text_width, preserve_paragraphs=True + if len(widths) < 3: + raise TypeError( + "Expected three columns for tagged definition list entries." + ) + + first_col = min(widths[0], col_max) + col_spacing + second_col = min(widths[1], col_max) + col_spacing * 2 + + current_tag: Optional[str] = None + + for first, second, third in iter_rows(rows, 3): + if current_tag != first: + current_tag = first + self.write("\n") + self.write( + f"[#431d93]{'':>{self.current_indent}}{first}:[/#431d93]\n" ) - lines = wrapped_text.splitlines() - - if lines: - self.write(f"{lines[0]}\n") - - for line in lines[1:]: - self.write( - f"{'':>{first_col + self.current_indent}}{line}\n" - ) - else: - self.write("\n") - - elif len(widths) == 3: - first_col = min(widths[0], col_max) + col_spacing - second_col = min(widths[1], col_max) + col_spacing * 2 - - current_tag = None - for first, second, third in iter_rows(rows, len(widths)): - if current_tag != first: - current_tag = first - self.write("\n") - # Adding [#431d93] [/#431d93] makes the tag colorful when - # it is printed by rich print - self.write( - f"[#431d93]{'':>{self.current_indent}}{first}:[/#431d93]\n" - ) - if not third: - self.write("\n") - continue + if not third: + self.write("\n") + continue + + if term_len(first) <= first_col - col_spacing: + self.write(" " * self.current_indent * 2) + else: + self.write("\n") + self.write(" " * (first_col + self.current_indent)) - if term_len(first) <= first_col - col_spacing: - self.write(" " * self.current_indent * 2) - else: - self.write("\n") - self.write(" " * (first_col + self.current_indent)) + self.write(f"{'':>{self.current_indent}}{second}") - self.write(f"{'':>{self.current_indent}}{second}") + text_width = max(self.width - second_col - 4, 10) + wrapped_text = formatting.wrap_text( + third, text_width, preserve_paragraphs=True + ) + lines = wrapped_text.splitlines() - text_width = max(self.width - second_col - 4, 10) - wrapped_text = formatting.wrap_text( - third, text_width, preserve_paragraphs=True + if lines: + self.write( + " " * (second_col - term_len(second) + self.current_indent) ) - lines = wrapped_text.splitlines() + self.write(f"{lines[0]}\n") - if lines: + for line in lines[1:]: self.write( - " " - * (second_col - term_len(second) + self.current_indent) + f"{'':>{second_col + self.current_indent * 4}}{line}\n" ) - self.write(f"{lines[0]}\n") - - for line in lines[1:]: - self.write( - f"{'':>{second_col + self.current_indent * 4}}{line}\n" - ) - else: - self.write("\n") - else: - raise TypeError( - "Expected either three or two columns for definition list" - ) + else: + self.write("\n") From d28f68b468d55da4689d7ff160e0c027e270b300 Mon Sep 17 00:00:00 2001 From: Alex Strick van Linschoten Date: Fri, 10 Oct 2025 17:15:55 +0200 Subject: [PATCH 02/25] Fix typo --- src/zenml/cli/formatter.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/zenml/cli/formatter.py b/src/zenml/cli/formatter.py index 3d0410d57ff..bf495e196ca 100644 --- a/src/zenml/cli/formatter.py +++ b/src/zenml/cli/formatter.py @@ -83,7 +83,7 @@ def write_dl( Click 8.2 tightened validation so that definition list entries must be pairs. Our CLI groups commands in tagged triples, so we detect that - case and render it manually while delegating the standard behaviour to + case and render it manually while delegating the standard behavior to Click for the classic two-column output. This keeps the formatter compatible with both older Click releases (<8.2) and the newer ones without sacrificing the custom grouped layout. @@ -98,7 +98,6 @@ def write_dl( TypeError: If the provided rows do not represent two or three column entries. """ - normalized_rows: List[Tuple[str, ...]] = [tuple(row) for row in rows] if not normalized_rows: From 5d07137c1a9f6e21e10042056fc5a02a71da25d8 Mon Sep 17 00:00:00 2001 From: Alex Strick van Linschoten Date: Mon, 20 Oct 2025 16:29:23 +0200 Subject: [PATCH 03/25] Fixes for CI --- src/zenml/cli/cli.py | 35 ++++++++++++++++++++++-------- src/zenml/cli/formatter.py | 44 ++++++++++++++++++++++++++++++++------ 2 files changed, 64 insertions(+), 15 deletions(-) diff --git a/src/zenml/cli/cli.py b/src/zenml/cli/cli.py index 01346581782..ab59eb3ff73 100644 --- a/src/zenml/cli/cli.py +++ b/src/zenml/cli/cli.py @@ -137,16 +137,24 @@ def format_commands( help_ = cmd.get_short_help_str(limit=formatter.width) rows.append((tag.value, subcommand, help_)) if rows: - colored_section_title = ( - "[dim cyan]Available ZenML Commands (grouped)[/dim cyan]" - ) - with formatter.section(colored_section_title): - formatter.write_dl(rows) # type: ignore[arg-type] - - -@click.group(cls=ZenMLCLI) + if isinstance(formatter, ZenFormatter): + section_title = "[dim cyan]Available ZenML Commands (grouped)[/dim cyan]" + with formatter.section(section_title): + formatter.write_dl(rows) # type: ignore[arg-type] + else: + # Fallback: use simple pairs without category and avoid rich markup in header + section_title = "Available ZenML Commands" + with formatter.section(section_title): + pair_rows: List[Tuple[str, str]] = [ + (subcmd, help_) for _, subcmd, help_ in rows + ] + formatter.write_dl(pair_rows) + + +@click.group(cls=ZenMLCLI, invoke_without_command=True) @click.version_option(__version__, "--version", "-v") -def cli() -> None: +@click.pass_context +def cli(ctx: click.Context) -> None: """CLI base command for ZenML.""" set_root_verbosity() source_context.set(SourceContextTypes.CLI) @@ -158,6 +166,15 @@ def cli() -> None: # packages directory source_utils.set_custom_source_root(source_root=os.getcwd()) + # Manually show help and exit with code 0 when invoked without a subcommand. + # Click 8.2+ raises NoArgsIsHelpError before the callback runs when + # no_args_is_help=True. By relying solely on invoke_without_command=True + # and handling help here, we ensure consistent behavior across Click + # versions while leveraging our custom help formatter in ZenMLCLI.get_help(). + if ctx.invoked_subcommand is None and not ctx.resilient_parsing: + ctx.command.get_help(ctx) + ctx.exit(0) + if __name__ == "__main__": cli() diff --git a/src/zenml/cli/formatter.py b/src/zenml/cli/formatter.py index bf495e196ca..2bc76d5dee9 100644 --- a/src/zenml/cli/formatter.py +++ b/src/zenml/cli/formatter.py @@ -16,7 +16,29 @@ from typing import Dict, Iterable, Iterator, List, Optional, Sequence, Tuple from click import formatting -from click._compat import term_len + +try: + # Prefer the public API available in newer Click versions + from click.utils import term_len # type: ignore[attr-defined] +except Exception: + # Fallback for older Click versions + from click._compat import term_len # type: ignore[attr-defined] + + +def _safe_width( + width: Optional[int], max_width: Optional[int], default: int = 80 +) -> int: + """Return a non-None width value suitable for wrapping. + + Click sets width to None in some non-interactive contexts. This helper + ensures we always get a valid integer width, preferring the explicit width + first, then max_width, then a sensible default. + """ + if isinstance(width, int) and width > 0: + return width + if isinstance(max_width, int) and max_width > 0: + return max_width + return default def measure_table(rows: Iterable[Tuple[str, ...]]) -> Tuple[int, ...]: @@ -137,12 +159,14 @@ def _write_triple_definition_list( "Expected three columns for tagged definition list entries." ) + # Compute target column offsets with spacing applied first_col = min(widths[0], col_max) + col_spacing second_col = min(widths[1], col_max) + col_spacing * 2 current_tag: Optional[str] = None for first, second, third in iter_rows(rows, 3): + # Print a category header when the tag changes if current_tag != first: current_tag = first self.write("\n") @@ -150,33 +174,41 @@ def _write_triple_definition_list( f"[#431d93]{'':>{self.current_indent}}{first}:[/#431d93]\n" ) + # Preserve behavior for empty third-column values if not third: self.write("\n") continue + # Layout for the command name column if term_len(first) <= first_col - col_spacing: - self.write(" " * self.current_indent * 2) + self.write(" " * (self.current_indent * 2)) else: self.write("\n") self.write(" " * (first_col + self.current_indent)) + # Command name with current indentation self.write(f"{'':>{self.current_indent}}{second}") - text_width = max(self.width - second_col - 4, 10) + # Wrap and render the description using a safe width calculation + effective_width = _safe_width( + self.width, getattr(self, "max_width", None) + ) + text_width = max(effective_width - second_col - 4, 10) wrapped_text = formatting.wrap_text( third, text_width, preserve_paragraphs=True ) lines = wrapped_text.splitlines() if lines: + # First line appears on the same row as the command name self.write( " " * (second_col - term_len(second) + self.current_indent) ) self.write(f"{lines[0]}\n") + # Continuation lines align with the description column for line in lines[1:]: - self.write( - f"{'':>{second_col + self.current_indent * 4}}{line}\n" - ) + self.write(" " * (second_col + self.current_indent)) + self.write(f"{line}\n") else: self.write("\n") From 5ce39d57f88672e4d530bd5cb8ac205b5d577146 Mon Sep 17 00:00:00 2001 From: Alex Strick van Linschoten Date: Mon, 20 Oct 2025 17:00:05 +0200 Subject: [PATCH 04/25] mypy fix --- src/zenml/cli/cli.py | 2 +- src/zenml/cli/formatter.py | 26 +++++++++++++++++++------- 2 files changed, 20 insertions(+), 8 deletions(-) diff --git a/src/zenml/cli/cli.py b/src/zenml/cli/cli.py index ab59eb3ff73..49c07e892db 100644 --- a/src/zenml/cli/cli.py +++ b/src/zenml/cli/cli.py @@ -140,7 +140,7 @@ def format_commands( if isinstance(formatter, ZenFormatter): section_title = "[dim cyan]Available ZenML Commands (grouped)[/dim cyan]" with formatter.section(section_title): - formatter.write_dl(rows) # type: ignore[arg-type] + formatter.write_dl(rows) else: # Fallback: use simple pairs without category and avoid rich markup in header section_title = "Available ZenML Commands" diff --git a/src/zenml/cli/formatter.py b/src/zenml/cli/formatter.py index 2bc76d5dee9..c5b0e6036fe 100644 --- a/src/zenml/cli/formatter.py +++ b/src/zenml/cli/formatter.py @@ -13,16 +13,28 @@ # permissions and limitations under the License. """Helper functions to format output for CLI.""" -from typing import Dict, Iterable, Iterator, List, Optional, Sequence, Tuple +from typing import ( + TYPE_CHECKING, + Dict, + Iterable, + Iterator, + List, + Optional, + Sequence, + Tuple, +) from click import formatting -try: - # Prefer the public API available in newer Click versions - from click.utils import term_len # type: ignore[attr-defined] -except Exception: - # Fallback for older Click versions - from click._compat import term_len # type: ignore[attr-defined] +if TYPE_CHECKING: + from typing import Callable + + term_len: Callable[[str], int] +else: + try: + from click.utils import term_len + except Exception: + from click._compat import term_len def _safe_width( From 747616845a4f1be7120b24d27b7606fa238d30c1 Mon Sep 17 00:00:00 2001 From: Alex Strick van Linschoten Date: Mon, 20 Oct 2025 18:03:54 +0200 Subject: [PATCH 05/25] Darglint fixes --- src/zenml/cli/cli.py | 6 +++++- src/zenml/cli/formatter.py | 9 +++++++++ 2 files changed, 14 insertions(+), 1 deletion(-) diff --git a/src/zenml/cli/cli.py b/src/zenml/cli/cli.py index 49c07e892db..20ef5f79d24 100644 --- a/src/zenml/cli/cli.py +++ b/src/zenml/cli/cli.py @@ -155,7 +155,11 @@ def format_commands( @click.version_option(__version__, "--version", "-v") @click.pass_context def cli(ctx: click.Context) -> None: - """CLI base command for ZenML.""" + """CLI base command for ZenML. + + Args: + ctx: The click context. + """ set_root_verbosity() source_context.set(SourceContextTypes.CLI) repo_root = Client.find_repository() diff --git a/src/zenml/cli/formatter.py b/src/zenml/cli/formatter.py index c5b0e6036fe..e578a9c3a11 100644 --- a/src/zenml/cli/formatter.py +++ b/src/zenml/cli/formatter.py @@ -45,6 +45,15 @@ def _safe_width( Click sets width to None in some non-interactive contexts. This helper ensures we always get a valid integer width, preferring the explicit width first, then max_width, then a sensible default. + + Args: + width: Explicit terminal or content width to use when available. + max_width: Upper bound to apply if width is not set or invalid. + default: Fallback width used when neither width nor max_width is a + positive integer. + + Returns: + The effective positive integer width used for wrapping. """ if isinstance(width, int) and width > 0: return width From db1907a4da2bee3822cc2df59fd22b37da617390 Mon Sep 17 00:00:00 2001 From: Alex Strick van Linschoten Date: Mon, 20 Oct 2025 18:05:34 +0200 Subject: [PATCH 06/25] Fix type hint for term_len in formatter.py Updated the type hint for term_len to ensure proper type checking during development. This change enhances code clarity and maintains compatibility with mypy checks. --- src/zenml/cli/formatter.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/zenml/cli/formatter.py b/src/zenml/cli/formatter.py index e578a9c3a11..1fe009bf60b 100644 --- a/src/zenml/cli/formatter.py +++ b/src/zenml/cli/formatter.py @@ -15,6 +15,7 @@ from typing import ( TYPE_CHECKING, + Callable, Dict, Iterable, Iterator, @@ -27,8 +28,6 @@ from click import formatting if TYPE_CHECKING: - from typing import Callable - term_len: Callable[[str], int] else: try: From ad309a80dbebb7ee38abc42692200e8c5c23c1ab Mon Sep 17 00:00:00 2001 From: Alex Strick van Linschoten Date: Mon, 20 Oct 2025 19:59:13 +0200 Subject: [PATCH 07/25] Fix CLI test failures with Click >=8.2 by using compatibility helper Tests were directly instantiating CliRunner() which fails with Click 8.2+ since the mix_stderr parameter was removed. The codebase already had a cli_runner() compatibility helper in utils.py that handles version differences, but tests weren't using it. This fix updates test_base.py and test_cli.py to use the helper function and adds exit code checks to catch failures early with clear error messages. Resolves template init test failures in CI where Click 8.2.1 is installed. --- tests/integration/functional/cli/test_base.py | 17 ++++++++++------- tests/integration/functional/cli/test_cli.py | 10 ++++++---- 2 files changed, 16 insertions(+), 11 deletions(-) diff --git a/tests/integration/functional/cli/test_base.py b/tests/integration/functional/cli/test_base.py index 7217a227712..46b0682920d 100644 --- a/tests/integration/functional/cli/test_base.py +++ b/tests/integration/functional/cli/test_base.py @@ -17,8 +17,8 @@ from pathlib import Path import pytest -from click.testing import CliRunner +from tests.integration.functional.cli.utils import cli_runner from zenml.cli.base import ZENML_PROJECT_TEMPLATES, clean, init from zenml.constants import CONFIG_FILE_NAME, REPOSITORY_DIRECTORY_NAME from zenml.utils import yaml_utils @@ -27,8 +27,9 @@ def test_init_creates_zen_folder(tmp_path: Path) -> None: """Check that init command creates a .zen folder.""" - runner = CliRunner() - runner.invoke(init, ["--path", str(tmp_path)]) + runner = cli_runner() + result = runner.invoke(init, ["--path", str(tmp_path)]) + assert result.exit_code == 0, f"Command failed: {result.output}" assert (tmp_path / REPOSITORY_DIRECTORY_NAME).exists() @@ -44,8 +45,8 @@ def test_init_creates_from_templates( tmp_path: Path, template_name: str ) -> None: """Check that init command checks-out template.""" - runner = CliRunner() - runner.invoke( + runner = cli_runner() + result = runner.invoke( init, [ "--path", @@ -55,6 +56,7 @@ def test_init_creates_from_templates( "--template-with-defaults", ], ) + assert result.exit_code == 0, f"Command failed: {result.output}" assert (tmp_path / REPOSITORY_DIRECTORY_NAME).exists() files_in_top_level = set([f.lower() for f in os.listdir(str(tmp_path))]) must_have_files = { @@ -75,8 +77,9 @@ def test_clean_user_config(clean_client) -> None: user_id = yaml_contents["user_id"] analytics_opt_in = yaml_contents["analytics_opt_in"] version = yaml_contents["version"] - runner = CliRunner() - runner.invoke(clean, ["--yes", True]) + runner = cli_runner() + result = runner.invoke(clean, ["--yes", True]) + assert result.exit_code == 0, f"Command failed: {result.output}" new_yaml_contents = yaml_utils.read_yaml(str(global_zen_config_yaml)) assert user_id == new_yaml_contents["user_id"] assert analytics_opt_in == new_yaml_contents["analytics_opt_in"] diff --git a/tests/integration/functional/cli/test_cli.py b/tests/integration/functional/cli/test_cli.py index 010f442c928..8da0c6d152d 100644 --- a/tests/integration/functional/cli/test_cli.py +++ b/tests/integration/functional/cli/test_cli.py @@ -16,15 +16,17 @@ import click import pytest -from click.testing import CliRunner +from tests.integration.functional.cli.utils import ( + cli_runner as create_cli_runner, +) from zenml.cli.cli import ZenMLCLI, cli from zenml.cli.formatter import ZenFormatter @pytest.fixture(scope="function") def runner(request): - return CliRunner() + return create_cli_runner() def test_cli_command_defines_a_cli_group() -> None: @@ -57,7 +59,7 @@ def test_cli_sets_custom_source_root_if_outside_of_repository( mock_set_custom_source_root = mocker.patch( "zenml.utils.source_utils.set_custom_source_root" ) - runner = CliRunner() + runner = create_cli_runner() # Invoke a subcommand so the root CLI group gets called runner.invoke(cli, ["version"]) @@ -73,7 +75,7 @@ def test_cli_does_not_set_custom_source_root_if_inside_repository( mock_set_custom_source_root = mocker.patch( "zenml.utils.source_utils.set_custom_source_root" ) - runner = CliRunner() + runner = create_cli_runner() # Invoke a subcommand so the root CLI group gets called runner.invoke(cli, ["version"]) From 515610f75c1c263643b537f4152e94599216c187 Mon Sep 17 00:00:00 2001 From: Alex Strick van Linschoten Date: Mon, 20 Oct 2025 20:07:14 +0200 Subject: [PATCH 08/25] Document CLI testing best practices for Click compatibility Add guidance for future CLI test development to use the cli_runner() helper for Click version compatibility, while clarifying that existing tests with plain CliRunner() calls are fine as-is. --- CLAUDE.md | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/CLAUDE.md b/CLAUDE.md index 1297d6d2838..8fbf744e74b 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -51,6 +51,12 @@ Use filesystem navigation tools to explore the codebase structure as needed. - For full coverage, use CI (see CI section below) - Some tests use: `bash scripts/test-coverage-xml.sh` (but this won't run all tests) +#### CLI Testing Best Practices +- When writing new CLI tests, use the `cli_runner()` helper from `tests/integration/functional/cli/utils.py` instead of directly instantiating `CliRunner()` +- The helper handles Click version compatibility (Click 8.2+ removed the `mix_stderr` parameter) +- Always check `result.exit_code` after invoking CLI commands to catch failures early +- Existing tests with `CliRunner()` (no arguments) are fine - only new tests need the helper + ## Development Workflow ### Prerequisites From 37ba025f7e2971cdd96c6709adef8d87447f8ea6 Mon Sep 17 00:00:00 2001 From: Alex Strick van Linschoten Date: Mon, 20 Oct 2025 21:34:34 +0200 Subject: [PATCH 09/25] Add --test flag to init command tests to bypass email prompt The init command prompts for email during initialization, which causes CI tests to fail with "Email:Aborted!" when running non-interactively. The --test flag (hidden CLI option) specifically bypasses this prompt for testing scenarios. This fixes test_init_creates_zen_folder and test_init_creates_from_templates which were failing in CI due to the interactive email prompt introduced when the command runs. --- tests/integration/functional/cli/test_base.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/integration/functional/cli/test_base.py b/tests/integration/functional/cli/test_base.py index 46b0682920d..f928024af4d 100644 --- a/tests/integration/functional/cli/test_base.py +++ b/tests/integration/functional/cli/test_base.py @@ -28,7 +28,7 @@ def test_init_creates_zen_folder(tmp_path: Path) -> None: """Check that init command creates a .zen folder.""" runner = cli_runner() - result = runner.invoke(init, ["--path", str(tmp_path)]) + result = runner.invoke(init, ["--path", str(tmp_path), "--test"]) assert result.exit_code == 0, f"Command failed: {result.output}" assert (tmp_path / REPOSITORY_DIRECTORY_NAME).exists() @@ -54,6 +54,7 @@ def test_init_creates_from_templates( "--template", template_name, "--template-with-defaults", + "--test", ], ) assert result.exit_code == 0, f"Command failed: {result.output}" From e331ce8a5d0a5f20af9c2793ad7d0d3c186ec096 Mon Sep 17 00:00:00 2001 From: Alex Strick van Linschoten Date: Tue, 21 Oct 2025 08:33:05 +0200 Subject: [PATCH 10/25] Improve exception handling and type hints in CLI formatter - Use ImportError instead of broad Exception when catching Click import failures This is more precise and only catches the expected failure case when term_len is unavailable in older Click versions, avoiding accidentally hiding other errors. - Add type: ignore[attr-defined] comment with explanation for max_width access Documents why we're bypassing type checking - max_width may not be present on all formatter types, which is intentional behavior. --- src/zenml/cli/formatter.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/zenml/cli/formatter.py b/src/zenml/cli/formatter.py index 1fe009bf60b..b98ce52f4f8 100644 --- a/src/zenml/cli/formatter.py +++ b/src/zenml/cli/formatter.py @@ -32,7 +32,7 @@ else: try: from click.utils import term_len - except Exception: + except ImportError: from click._compat import term_len @@ -211,7 +211,8 @@ def _write_triple_definition_list( # Wrap and render the description using a safe width calculation effective_width = _safe_width( - self.width, getattr(self, "max_width", None) + self.width, + getattr(self, "max_width", None), # type: ignore[attr-defined] # max_width may not be present on all formatter types; this is intentional ) text_width = max(effective_width - second_col - 4, 10) wrapped_text = formatting.wrap_text( From 92baa4d379be5f59000e8c06dd5c6741b0b4502f Mon Sep 17 00:00:00 2001 From: Alex Strick van Linschoten Date: Tue, 21 Oct 2025 08:41:44 +0200 Subject: [PATCH 11/25] Fix clean command test to properly use boolean flag The --yes flag is defined as is_flag=True in Click, which means it's a boolean flag that works by presence/absence, not by value assignment. Changed from ["--yes", True] to ["--yes"] to fix the TypeError where Click was trying to process the boolean True as a string argument. --- tests/integration/functional/cli/test_base.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/integration/functional/cli/test_base.py b/tests/integration/functional/cli/test_base.py index f928024af4d..4b5b7df692c 100644 --- a/tests/integration/functional/cli/test_base.py +++ b/tests/integration/functional/cli/test_base.py @@ -79,7 +79,7 @@ def test_clean_user_config(clean_client) -> None: analytics_opt_in = yaml_contents["analytics_opt_in"] version = yaml_contents["version"] runner = cli_runner() - result = runner.invoke(clean, ["--yes", True]) + result = runner.invoke(clean, ["--yes"]) assert result.exit_code == 0, f"Command failed: {result.output}" new_yaml_contents = yaml_utils.read_yaml(str(global_zen_config_yaml)) assert user_id == new_yaml_contents["user_id"] From 5fae5633b5f24395a4c6fbd8e521497e777d208e Mon Sep 17 00:00:00 2001 From: Alex Strick van Linschoten Date: Tue, 21 Oct 2025 08:50:52 +0200 Subject: [PATCH 12/25] Remove unused type: ignore comment in CLI formatter Mypy no longer requires the attr-defined suppression for the getattr call on max_width. The explanatory comment remains to document the intentional use of getattr for compatibility. --- src/zenml/cli/formatter.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/zenml/cli/formatter.py b/src/zenml/cli/formatter.py index b98ce52f4f8..d66e2ec084d 100644 --- a/src/zenml/cli/formatter.py +++ b/src/zenml/cli/formatter.py @@ -212,7 +212,9 @@ def _write_triple_definition_list( # Wrap and render the description using a safe width calculation effective_width = _safe_width( self.width, - getattr(self, "max_width", None), # type: ignore[attr-defined] # max_width may not be present on all formatter types; this is intentional + getattr( + self, "max_width", None + ), # max_width may not be present on all formatter types; this is intentional ) text_width = max(effective_width - second_col - 4, 10) wrapped_text = formatting.wrap_text( From 65abb21b67fd15b95cc0442c8a2ccdae8d6360fa Mon Sep 17 00:00:00 2001 From: Alex Strick van Linschoten Date: Tue, 21 Oct 2025 10:34:57 +0200 Subject: [PATCH 13/25] Add missing init file --- tests/integration/functional/utils/__init__.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) create mode 100644 tests/integration/functional/utils/__init__.py diff --git a/tests/integration/functional/utils/__init__.py b/tests/integration/functional/utils/__init__.py new file mode 100644 index 00000000000..32ba63d3703 --- /dev/null +++ b/tests/integration/functional/utils/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) ZenML GmbH 2025. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at: +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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. From 0ba7fdf2971db7befd70862e0a03a34781bb09cb Mon Sep 17 00:00:00 2001 From: Alex Strick van Linschoten Date: Tue, 21 Oct 2025 10:35:21 +0200 Subject: [PATCH 14/25] Formatting --- tests/integration/functional/utils/test_tag_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/integration/functional/utils/test_tag_utils.py b/tests/integration/functional/utils/test_tag_utils.py index fe684372b58..be5329a4e69 100644 --- a/tests/integration/functional/utils/test_tag_utils.py +++ b/tests/integration/functional/utils/test_tag_utils.py @@ -17,8 +17,8 @@ from typing import Annotated, Tuple import pytest -from tests.integration.functional.zen_stores.utils import PipelineRunContext +from tests.integration.functional.zen_stores.utils import PipelineRunContext from zenml import ArtifactConfig, Tag, add_tags, pipeline, remove_tags, step from zenml.client import Client from zenml.constants import ENV_ZENML_PREVENT_CLIENT_SIDE_CACHING From 8cca137b7c718ce9bbb135a807f9450b19f0b1ca Mon Sep 17 00:00:00 2001 From: Alex Strick van Linschoten Date: Tue, 21 Oct 2025 10:55:47 +0200 Subject: [PATCH 15/25] Fix test failures by renaming utils test directory Previously, tests in tests/integration/functional/utils/ were failing during pipeline execution with ModuleNotFoundError. The directory was missing __init__.py, preventing ZenML from importing step definitions. Adding __init__.py created a namespace collision between the utils.py module (exporting sample_name) and the utils/ package directory. This broke imports across multiple test files that depend on sample_name. Renamed utils/ to utils_functional/ to eliminate the shadowing issue while preserving the utils.py module for shared test utilities. --- .../functional/{utils => utils_functional}/__init__.py | 0 .../functional/{utils => utils_functional}/test_metadata_utils.py | 0 .../functional/{utils => utils_functional}/test_run_utils.py | 0 .../functional/{utils => utils_functional}/test_tag_utils.py | 0 4 files changed, 0 insertions(+), 0 deletions(-) rename tests/integration/functional/{utils => utils_functional}/__init__.py (100%) rename tests/integration/functional/{utils => utils_functional}/test_metadata_utils.py (100%) rename tests/integration/functional/{utils => utils_functional}/test_run_utils.py (100%) rename tests/integration/functional/{utils => utils_functional}/test_tag_utils.py (100%) diff --git a/tests/integration/functional/utils/__init__.py b/tests/integration/functional/utils_functional/__init__.py similarity index 100% rename from tests/integration/functional/utils/__init__.py rename to tests/integration/functional/utils_functional/__init__.py diff --git a/tests/integration/functional/utils/test_metadata_utils.py b/tests/integration/functional/utils_functional/test_metadata_utils.py similarity index 100% rename from tests/integration/functional/utils/test_metadata_utils.py rename to tests/integration/functional/utils_functional/test_metadata_utils.py diff --git a/tests/integration/functional/utils/test_run_utils.py b/tests/integration/functional/utils_functional/test_run_utils.py similarity index 100% rename from tests/integration/functional/utils/test_run_utils.py rename to tests/integration/functional/utils_functional/test_run_utils.py diff --git a/tests/integration/functional/utils/test_tag_utils.py b/tests/integration/functional/utils_functional/test_tag_utils.py similarity index 100% rename from tests/integration/functional/utils/test_tag_utils.py rename to tests/integration/functional/utils_functional/test_tag_utils.py From be76ce525f416b07401ea8ab235c1641fdf589f8 Mon Sep 17 00:00:00 2001 From: Alex Strick van Linschoten Date: Tue, 21 Oct 2025 13:36:19 +0200 Subject: [PATCH 16/25] Fix test collection failures in test_step_context Pytest collection was failing due to import-time Client and materializer resolution before fixtures could configure the test environment. Changes: - Move Client import under TYPE_CHECKING guard to prevent import-time initialization - Remove output_materializers from @step decorator to avoid import-time resolution - Apply materializers at runtime using .with_options() inside test functions - Add clean_client fixture to all pipeline-running tests for proper setup This ensures materializers can access StepContext during save/load operations without triggering "No active project configured" errors during collection. --- .../functional/steps/test_step_context.py | 22 +++++++++---------- 1 file changed, 11 insertions(+), 11 deletions(-) diff --git a/tests/integration/functional/steps/test_step_context.py b/tests/integration/functional/steps/test_step_context.py index f34e53b8e4b..dc258255bbf 100644 --- a/tests/integration/functional/steps/test_step_context.py +++ b/tests/integration/functional/steps/test_step_context.py @@ -1,6 +1,6 @@ import json import os -from typing import Any, Dict, List, Optional, Tuple, Type +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Type import pytest from pydantic import BaseModel @@ -8,11 +8,13 @@ from zenml import get_step_context, log_metadata, pipeline, step from zenml.artifacts.artifact_config import ArtifactConfig -from zenml.client import Client from zenml.enums import ArtifactType from zenml.io import fileio from zenml.materializers.base_materializer import BaseMaterializer +if TYPE_CHECKING: + from zenml.client import Client + class ComplexObject(BaseModel): """This is a custom object to be materialized with ComplexObjectMaterializer.""" @@ -44,11 +46,7 @@ def save(self, data: ComplexObject) -> None: f.write(data.model_dump_json()) -@step( - output_materializers=[ - ComplexObjectMaterializer, - ] -) +@step def _output_complex_object_step(): """This step would call `save` of `ComplexObjectMaterializer`. `save` should not fail and have access to the step context""" @@ -68,19 +66,21 @@ def _access_step_context_step(): assert get_step_context().pipeline.name == "bar" -def test_materializer_can_access_step_context(): +def test_materializer_can_access_step_context(clean_client: "Client"): """Call steps using `ComplexObjectMaterializer` to validate that step context is available to Materializers""" @pipeline(name="bar") def _complex_object_materialization_pipeline(): - complex_object = _output_complex_object_step() + complex_object = _output_complex_object_step.with_options( + output_materializers=ComplexObjectMaterializer + )() _input_complex_object_step(complex_object) _complex_object_materialization_pipeline() -def test_step_can_access_step_context(): +def test_step_can_access_step_context(clean_client: "Client"): """Call step using step context directly, before Materializers""" @pipeline(name="bar") @@ -106,7 +106,7 @@ def step_context_metadata_reader_step(my_input: int) -> None: assert my_input_metadata["some_key"] == "some_value" -def test_input_artifacts_property(): +def test_input_artifacts_property(clean_client: "Client"): """Test the `StepContext.inputs` property.""" @pipeline From 60f14ae9aef7afeb99cc2a22fc4f7253342fc566 Mon Sep 17 00:00:00 2001 From: Alex Strick van Linschoten Date: Tue, 21 Oct 2025 17:58:19 +0200 Subject: [PATCH 17/25] Fix failing windows DB lock in integration tests --- src/zenml/cli/base.py | 2 + src/zenml/config/global_config.py | 35 +++++++++++++ src/zenml/zen_stores/sql_zen_store.py | 52 +++++++++++++++++++ tests/integration/functional/cli/test_base.py | 30 +++++++++++ 4 files changed, 119 insertions(+) diff --git a/src/zenml/cli/base.py b/src/zenml/cli/base.py index d80f428b92d..5256e49d0da 100644 --- a/src/zenml/cli/base.py +++ b/src/zenml/cli/base.py @@ -381,6 +381,8 @@ def clean(yes: bool = False, local: bool = False) -> None: global_zen_config = Path(get_global_config_directory()) if fileio.exists(str(global_zen_config)): gc = GlobalConfiguration() + # Ensure any open DB handles (e.g., SQLite) are released before deletion + gc.close_store() for dir_name in fileio.listdir(str(global_zen_config)): if fileio.isdir(str(global_zen_config / str(dir_name))): warning( diff --git a/src/zenml/config/global_config.py b/src/zenml/config/global_config.py index 4cc8d17fb3b..c46a926f8e8 100644 --- a/src/zenml/config/global_config.py +++ b/src/zenml/config/global_config.py @@ -170,6 +170,17 @@ def _reset_instance( singleton. If None, the global GlobalConfiguration singleton is reset to an empty value. """ + # Proactively dispose resources held by the current singleton instance + # to avoid file-handle issues on platforms like Windows. + if cls._global_config is not None: + try: + cls._global_config.close_store() + except Exception: + logger.debug( + "Failed to close store during GlobalConfiguration reset.", + exc_info=True, + ) + cls._global_config = config if config: config._write_config() @@ -701,6 +712,30 @@ def set_store( self._configure_store(config, skip_default_registrations, **kwargs) logger.info("Updated the global store configuration.") + def close_store(self) -> None: + """Close and detach the cached store if it has been initialized. + + Important: + - Operates only on the cached store; does not instantiate a new one. + - Logs cleanup errors at debug level and never raises. + """ + store = self._zen_store + if store is None: + return + + close_fn = getattr(store, "close", None) + if callable(close_fn): + try: + close_fn() + except Exception as e: + logger.debug( + "Error while closing zen store during cleanup: %s", + e, + exc_info=True, + ) + # Clear the cached store reference irrespective of cleanup result. + self._zen_store = None + @property def is_initialized(self) -> bool: """Check if the global configuration is initialized. diff --git a/src/zenml/zen_stores/sql_zen_store.py b/src/zenml/zen_stores/sql_zen_store.py index 3cf0af8b6dd..2faf3eceb1e 100644 --- a/src/zenml/zen_stores/sql_zen_store.py +++ b/src/zenml/zen_stores/sql_zen_store.py @@ -1238,6 +1238,30 @@ def filter_and_paginate( # Initialization and configuration # -------------------------------- + def close(self) -> None: + """Release DB resources held by this store. + + Safe to call multiple times. Primarily ensures SQLite file handles are + released so the DB file can be removed on platforms like Windows. + """ + try: + engine = self._engine + if engine is not None: + # Disposes pooled connections and releases file handles. + engine.dispose() + except Exception as e: + # Cleanup must not raise; surface at debug level only. + logger.debug( + "Error disposing SQLAlchemy engine during SqlZenStore.close(): %s", + e, + exc_info=True, + ) + finally: + # Clear references to avoid accidental reuse and help GC. + self._engine = None + self._alembic = None + self._migration_utils = None + def _initialize(self) -> None: """Initialize the SQL store.""" logger.debug("Initializing SqlZenStore at %s", self.config.url) @@ -1332,6 +1356,34 @@ def _initialize_database(self) -> None: # Send user enriched events that we missed due to a bug in 0.57.0 self._send_user_enriched_events_if_necessary() + # -------------------------------- + # Resource lifecycle + # -------------------------------- + + def close(self) -> None: + """Release DB resources held by this store. + + Safe to call multiple times. Primarily ensures SQLite file handles are + released so the DB file can be removed on platforms like Windows. + """ + try: + engine = self._engine + if engine is not None: + # Disposes pooled connections and releases file handles. + engine.dispose() + except Exception as e: + # Cleanup must not raise; surface at debug level only. + logger.debug( + "Error disposing SQLAlchemy engine during SqlZenStore.close(): %s", + e, + exc_info=True, + ) + finally: + # Clear references to avoid accidental reuse and help GC. + self._engine = None + self._alembic = None + self._migration_utils = None + def _get_db_backup_file_path(self) -> str: """Get the path to the database backup file. diff --git a/tests/integration/functional/cli/test_base.py b/tests/integration/functional/cli/test_base.py index 4b5b7df692c..84fb90669ee 100644 --- a/tests/integration/functional/cli/test_base.py +++ b/tests/integration/functional/cli/test_base.py @@ -85,3 +85,33 @@ def test_clean_user_config(clean_client) -> None: assert user_id == new_yaml_contents["user_id"] assert analytics_opt_in == new_yaml_contents["analytics_opt_in"] assert version == new_yaml_contents["version"] + + +def test_clean_calls_store_close(monkeypatch, clean_client) -> None: + # Import locally to avoid modifying module-level imports + from zenml.config.global_config import GlobalConfiguration + from zenml.zen_stores.sql_zen_store import SqlZenStore + + # Track whether close was called + called = {"closed": False} + + # Keep reference to the original close to preserve behavior + orig_close = SqlZenStore.close + + def tracked_close(self): + called["closed"] = True + # Call through to the original to keep cleanup behavior intact + return orig_close(self) + + # Patch the close method + monkeypatch.setattr(SqlZenStore, "close", tracked_close, raising=True) + + # Ensure the store is initialized in this process so there is something to close + _ = GlobalConfiguration().zen_store + + # Run clean + result = cli_runner().invoke(clean, ["--yes"]) + assert result.exit_code == 0, f"Command failed: {result.output}" + + # Verify that our close hook was executed + assert called["closed"] is True From 6dd8121d160564727e081f60abadc8ae1d2932e5 Mon Sep 17 00:00:00 2001 From: Alex Strick van Linschoten Date: Tue, 21 Oct 2025 18:10:07 +0200 Subject: [PATCH 18/25] Remove duplicated close method --- src/zenml/zen_stores/sql_zen_store.py | 28 --------------------------- 1 file changed, 28 deletions(-) diff --git a/src/zenml/zen_stores/sql_zen_store.py b/src/zenml/zen_stores/sql_zen_store.py index 2faf3eceb1e..685fbec17e1 100644 --- a/src/zenml/zen_stores/sql_zen_store.py +++ b/src/zenml/zen_stores/sql_zen_store.py @@ -1356,34 +1356,6 @@ def _initialize_database(self) -> None: # Send user enriched events that we missed due to a bug in 0.57.0 self._send_user_enriched_events_if_necessary() - # -------------------------------- - # Resource lifecycle - # -------------------------------- - - def close(self) -> None: - """Release DB resources held by this store. - - Safe to call multiple times. Primarily ensures SQLite file handles are - released so the DB file can be removed on platforms like Windows. - """ - try: - engine = self._engine - if engine is not None: - # Disposes pooled connections and releases file handles. - engine.dispose() - except Exception as e: - # Cleanup must not raise; surface at debug level only. - logger.debug( - "Error disposing SQLAlchemy engine during SqlZenStore.close(): %s", - e, - exc_info=True, - ) - finally: - # Clear references to avoid accidental reuse and help GC. - self._engine = None - self._alembic = None - self._migration_utils = None - def _get_db_backup_file_path(self) -> str: """Get the path to the database backup file. From de20afff310e40c5c41c646e17b21ada7ff5f709 Mon Sep 17 00:00:00 2001 From: Alex Strick van Linschoten Date: Wed, 22 Oct 2025 08:10:16 +0200 Subject: [PATCH 19/25] Define the close method on the ZenStoreInterface --- src/zenml/config/global_config.py | 23 +++++++------- src/zenml/zen_stores/base_zen_store.py | 12 ++++++++ src/zenml/zen_stores/rest_zen_store.py | 34 +++++++++++++++++++++ src/zenml/zen_stores/zen_store_interface.py | 13 ++++++++ 4 files changed, 70 insertions(+), 12 deletions(-) diff --git a/src/zenml/config/global_config.py b/src/zenml/config/global_config.py index c46a926f8e8..d764d7e0bf0 100644 --- a/src/zenml/config/global_config.py +++ b/src/zenml/config/global_config.py @@ -723,18 +723,17 @@ def close_store(self) -> None: if store is None: return - close_fn = getattr(store, "close", None) - if callable(close_fn): - try: - close_fn() - except Exception as e: - logger.debug( - "Error while closing zen store during cleanup: %s", - e, - exc_info=True, - ) - # Clear the cached store reference irrespective of cleanup result. - self._zen_store = None + try: + store.close() + except Exception as e: + logger.debug( + "Error while closing zen store during cleanup: %s", + e, + exc_info=True, + ) + finally: + # Clear the cached store reference irrespective of cleanup result. + self._zen_store = None @property def is_initialized(self) -> bool: diff --git a/src/zenml/zen_stores/base_zen_store.py b/src/zenml/zen_stores/base_zen_store.py index 1f369a6e817..614993a1115 100644 --- a/src/zenml/zen_stores/base_zen_store.py +++ b/src/zenml/zen_stores/base_zen_store.py @@ -141,6 +141,18 @@ def __init__( else: logger.debug("Skipping database initialization") + def close(self) -> None: + """Release external resources held by this store instance. + + Default no-op implementation. Subclasses should override this to + release any external resources (e.g., database connections, HTTP + sessions, file handles, background workers). Implementations must + remain idempotent (safe to call multiple times) and must not raise; + any cleanup failures should be handled internally and logged at + debug level instead. + """ + return + @staticmethod def get_store_class(store_type: StoreType) -> Type["BaseZenStore"]: """Returns the class of the given store type. diff --git a/src/zenml/zen_stores/rest_zen_store.py b/src/zenml/zen_stores/rest_zen_store.py index fdde0b41302..cefd724b072 100644 --- a/src/zenml/zen_stores/rest_zen_store.py +++ b/src/zenml/zen_stores/rest_zen_store.py @@ -4265,6 +4265,40 @@ def batch_delete_tag_resource( route=TAG_RESOURCES, ) + def close(self) -> None: + """Release external resources held by this store instance. + + This implementation closes the underlying HTTP session and clears transient + authentication state. It is safe to call multiple times and never raises; + unexpected cleanup issues are logged at debug level instead. + """ + try: + with self._session_lock: + if self._session is not None: + try: + self._session.close() + except Exception as e: + # Cleanup must not raise; surface at debug level only. + logger.debug( + "Error closing REST session during RestZenStore.close(): %s", + e, + exc_info=True, + ) + finally: + # Clear the session reference to avoid reuse and help GC. + self._session = None + + # Clear transient auth state so future use starts cleanly. + self._api_token = None + self._last_authenticated = None + except Exception as e: + # Absolute safety: no exceptions should escape close() + logger.debug( + "Unexpected error during RestZenStore.close(): %s", + e, + exc_info=True, + ) + # ======================= # Internal helper methods # ======================= diff --git a/src/zenml/zen_stores/zen_store_interface.py b/src/zenml/zen_stores/zen_store_interface.py index 8848a24b5ad..a975eae1867 100644 --- a/src/zenml/zen_stores/zen_store_interface.py +++ b/src/zenml/zen_stores/zen_store_interface.py @@ -231,6 +231,19 @@ def _initialize(self) -> None: be used to set up the backend (database, connection etc.). """ + @abstractmethod + def close(self) -> None: + """Release external resources held by this store instance. + + Implementations must: + - Release any external resources such as database connections, HTTP + sessions, open file handles, background workers, etc. + - Be safe to call multiple times (idempotent). Subsequent calls should + be no-ops and must not fail because resources were already released. + - Never raise exceptions. Any cleanup failures should be handled + internally and logged at debug level instead of being propagated. + """ + @abstractmethod def get_store_info(self) -> ServerModel: """Get information about the store. From 37dd2609df60d324a0098e4398129f52ad2f0e7d Mon Sep 17 00:00:00 2001 From: Alex Strick van Linschoten Date: Wed, 22 Oct 2025 08:18:17 +0200 Subject: [PATCH 20/25] Improve comment clarity on Click 8.2+ help behavior Clarified that the issue with Click 8.2+ isn't that help won't display, but that it would use Click's default formatter instead of ZenML's custom formatter, losing the custom categories, colors, and layout. --- src/zenml/cli/cli.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/src/zenml/cli/cli.py b/src/zenml/cli/cli.py index 20ef5f79d24..5cb692750cd 100644 --- a/src/zenml/cli/cli.py +++ b/src/zenml/cli/cli.py @@ -170,11 +170,13 @@ def cli(ctx: click.Context) -> None: # packages directory source_utils.set_custom_source_root(source_root=os.getcwd()) - # Manually show help and exit with code 0 when invoked without a subcommand. - # Click 8.2+ raises NoArgsIsHelpError before the callback runs when - # no_args_is_help=True. By relying solely on invoke_without_command=True - # and handling help here, we ensure consistent behavior across Click - # versions while leveraging our custom help formatter in ZenMLCLI.get_help(). + # Manually show help when invoked without a subcommand to ensure our custom + # formatter is used. Without invoke_without_command=True, Click defaults to + # no_args_is_help=True, which in Click 8.2+ raises NoArgsIsHelpError before + # this callback runs. That error triggers Click's default help display, + # bypassing ZenMLCLI.get_help() and losing our custom formatting (categories, + # colors, layout). By using invoke_without_command=True and manually checking + # for no subcommand, we guarantee consistent behavior across Click versions. if ctx.invoked_subcommand is None and not ctx.resilient_parsing: ctx.command.get_help(ctx) ctx.exit(0) From ff1056136c4d21b4fd5e993184b2db144cf7e958 Mon Sep 17 00:00:00 2001 From: Alex Strick van Linschoten Date: Wed, 22 Oct 2025 08:21:41 +0200 Subject: [PATCH 21/25] Remove unnecessary formatter fallback branch The isinstance check for ZenFormatter was overly defensive. Since ZenMLCLI defines context_class = ZenContext and ZenContext defines formatter_class = ZenFormatter, the formatter passed to format_commands() is guaranteed to be a ZenFormatter instance through Click's class hierarchy. The fallback branch was unreachable in practice. --- src/zenml/cli/cli.py | 17 +++++------------ 1 file changed, 5 insertions(+), 12 deletions(-) diff --git a/src/zenml/cli/cli.py b/src/zenml/cli/cli.py index 5cb692750cd..1e1a5c56f29 100644 --- a/src/zenml/cli/cli.py +++ b/src/zenml/cli/cli.py @@ -137,18 +137,11 @@ def format_commands( help_ = cmd.get_short_help_str(limit=formatter.width) rows.append((tag.value, subcommand, help_)) if rows: - if isinstance(formatter, ZenFormatter): - section_title = "[dim cyan]Available ZenML Commands (grouped)[/dim cyan]" - with formatter.section(section_title): - formatter.write_dl(rows) - else: - # Fallback: use simple pairs without category and avoid rich markup in header - section_title = "Available ZenML Commands" - with formatter.section(section_title): - pair_rows: List[Tuple[str, str]] = [ - (subcmd, help_) for _, subcmd, help_ in rows - ] - formatter.write_dl(pair_rows) + section_title = ( + "[dim cyan]Available ZenML Commands (grouped)[/dim cyan]" + ) + with formatter.section(section_title): + formatter.write_dl(rows) @click.group(cls=ZenMLCLI, invoke_without_command=True) From d9e61584437dd3986401ea4eabefa1ad66a1f608 Mon Sep 17 00:00:00 2001 From: Alex Strick van Linschoten Date: Wed, 22 Oct 2025 08:31:37 +0200 Subject: [PATCH 22/25] .gitignore updates --- .gitignore | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/.gitignore b/.gitignore index 266e685801f..58af9636a5c 100644 --- a/.gitignore +++ b/.gitignore @@ -195,9 +195,10 @@ examples/feast_feature_store/feast_feature_repo/data/registry.db #generated folder by zenml zenml_tutorial/ -# Claude settings +# AI code tools **/.claude/settings.local.json .claude/ +.repoprompt/ .local/ # PLEASE KEEP THIS LINE AT THE EOF: never include here src/zenml/zen_server/dashboard, since it is affecting release flow From aac47c2797c68813622c9717dc38891dbb548e37 Mon Sep 17 00:00:00 2001 From: Alex Strick van Linschoten Date: Wed, 22 Oct 2025 10:31:53 +0200 Subject: [PATCH 23/25] Simplify Click 8.2+ fix by reverting unnecessary formatter changes After review feedback, the formatter refactoring was not needed to fix Click 8.2+ compatibility. The formatter works fine on both old and new Click versions. The minimal fix required is: 1. Tests use cli_runner() helper (already in place) 2. CLI uses invoke_without_command=True for consistent help display (already in place) This commit reverts the formatter.py changes to keep the PR focused on the actual issue. --- src/zenml/cli/formatter.py | 244 +++++++++++++++---------------------- 1 file changed, 95 insertions(+), 149 deletions(-) diff --git a/src/zenml/cli/formatter.py b/src/zenml/cli/formatter.py index d66e2ec084d..1d35c757d36 100644 --- a/src/zenml/cli/formatter.py +++ b/src/zenml/cli/formatter.py @@ -13,52 +13,10 @@ # permissions and limitations under the License. """Helper functions to format output for CLI.""" -from typing import ( - TYPE_CHECKING, - Callable, - Dict, - Iterable, - Iterator, - List, - Optional, - Sequence, - Tuple, -) +from typing import Dict, Iterable, Iterator, Optional, Sequence, Tuple from click import formatting - -if TYPE_CHECKING: - term_len: Callable[[str], int] -else: - try: - from click.utils import term_len - except ImportError: - from click._compat import term_len - - -def _safe_width( - width: Optional[int], max_width: Optional[int], default: int = 80 -) -> int: - """Return a non-None width value suitable for wrapping. - - Click sets width to None in some non-interactive contexts. This helper - ensures we always get a valid integer width, preferring the explicit width - first, then max_width, then a sensible default. - - Args: - width: Explicit terminal or content width to use when available. - max_width: Upper bound to apply if width is not set or invalid. - default: Fallback width used when neither width nor max_width is a - positive integer. - - Returns: - The effective positive integer width used for wrapping. - """ - if isinstance(width, int) and width > 0: - return width - if isinstance(max_width, int) and max_width > 0: - return max_width - return default +from click._compat import term_len def measure_table(rows: Iterable[Tuple[str, ...]]) -> Tuple[int, ...]: @@ -121,117 +79,105 @@ def write_dl( col_max: int = 30, col_spacing: int = 2, ) -> None: - """Writes a definition list into the formatter buffer. + """Writes a definition list into the buffer. - Click 8.2 tightened validation so that definition list entries must be - pairs. Our CLI groups commands in tagged triples, so we detect that - case and render it manually while delegating the standard behavior to - Click for the classic two-column output. This keeps the formatter - compatible with both older Click releases (<8.2) and the newer ones - without sacrificing the custom grouped layout. + This is how options and commands are usually formatted. - Args: - rows: A sequence of tuples that represent the definition list - entries. - col_max: The maximum width of the first column. - col_spacing: The number of spaces between columns. + Arguments: + rows: a list of items as tuples for the terms and values. + col_max: the maximum width of the first column. + col_spacing: the number of spaces between the first and + second column (and third). + + The default behavior is to format the rows in a definition list + with rows of 2 columns following the format ``(term, value)``. + But for new CLI commands, we want to format the rows in a definition + list with rows of 3 columns following the format + ``(term, value, description)``. Raises: - TypeError: If the provided rows do not represent two or three - column entries. + TypeError: if the number of columns is not 2 or 3. """ - normalized_rows: List[Tuple[str, ...]] = [tuple(row) for row in rows] - - if not normalized_rows: - return - - unique_lengths = {len(row) for row in normalized_rows} - - if unique_lengths == {2}: - two_col_rows = [(row[0], row[1]) for row in normalized_rows] - super().write_dl( - two_col_rows, col_max=col_max, col_spacing=col_spacing - ) - return - - if unique_lengths == {3}: - self._write_triple_definition_list( - [(row[0], row[1], row[2]) for row in normalized_rows], - col_max=col_max, - col_spacing=col_spacing, - ) - return - - raise TypeError( - "Expected either two- or three-column definition list entries." - ) - - def _write_triple_definition_list( - self, - rows: List[Tuple[str, str, str]], - col_max: int, - col_spacing: int, - ) -> None: + rows = list(rows) widths = measure_table(rows) - if len(widths) < 3: - raise TypeError( - "Expected three columns for tagged definition list entries." - ) - - # Compute target column offsets with spacing applied - first_col = min(widths[0], col_max) + col_spacing - second_col = min(widths[1], col_max) + col_spacing * 2 - - current_tag: Optional[str] = None - - for first, second, third in iter_rows(rows, 3): - # Print a category header when the tag changes - if current_tag != first: - current_tag = first - self.write("\n") - self.write( - f"[#431d93]{'':>{self.current_indent}}{first}:[/#431d93]\n" + if len(widths) == 2: + first_col = min(widths[0], col_max) + col_spacing + + for first, second in iter_rows(rows, len(widths)): + self.write(f"{'':>{self.current_indent}}{first}") + if not second: + self.write("\n") + continue + if term_len(first) <= first_col - col_spacing: + self.write(" " * (first_col - term_len(first))) + else: + self.write("\n") + self.write(" " * (first_col + self.current_indent)) + + text_width = max(self.width - first_col - 2, 10) + wrapped_text = formatting.wrap_text( + second, text_width, preserve_paragraphs=True ) - - # Preserve behavior for empty third-column values - if not third: - self.write("\n") - continue - - # Layout for the command name column - if term_len(first) <= first_col - col_spacing: - self.write(" " * (self.current_indent * 2)) - else: - self.write("\n") - self.write(" " * (first_col + self.current_indent)) - - # Command name with current indentation - self.write(f"{'':>{self.current_indent}}{second}") - - # Wrap and render the description using a safe width calculation - effective_width = _safe_width( - self.width, - getattr( - self, "max_width", None - ), # max_width may not be present on all formatter types; this is intentional - ) - text_width = max(effective_width - second_col - 4, 10) - wrapped_text = formatting.wrap_text( - third, text_width, preserve_paragraphs=True - ) - lines = wrapped_text.splitlines() - - if lines: - # First line appears on the same row as the command name - self.write( - " " * (second_col - term_len(second) + self.current_indent) + lines = wrapped_text.splitlines() + + if lines: + self.write(f"{lines[0]}\n") + + for line in lines[1:]: + self.write( + f"{'':>{first_col + self.current_indent}}{line}\n" + ) + else: + self.write("\n") + + elif len(widths) == 3: + first_col = min(widths[0], col_max) + col_spacing + second_col = min(widths[1], col_max) + col_spacing * 2 + + current_tag = None + for first, second, third in iter_rows(rows, len(widths)): + if current_tag != first: + current_tag = first + self.write("\n") + # Adding [#431d93] [/#431d93] makes the tag colorful when + # it is printed by rich print + self.write( + f"[#431d93]{'':>{self.current_indent}}{first}:[/#431d93]\n" + ) + + if not third: + self.write("\n") + continue + + if term_len(first) <= first_col - col_spacing: + self.write(" " * self.current_indent * 2) + else: + self.write("\n") + self.write(" " * (first_col + self.current_indent)) + + self.write(f"{'':>{self.current_indent}}{second}") + + text_width = max(self.width - second_col - 4, 10) + wrapped_text = formatting.wrap_text( + third, text_width, preserve_paragraphs=True ) - self.write(f"{lines[0]}\n") - - # Continuation lines align with the description column - for line in lines[1:]: - self.write(" " * (second_col + self.current_indent)) - self.write(f"{line}\n") - else: - self.write("\n") + lines = wrapped_text.splitlines() + + if lines: + self.write( + " " + * (second_col - term_len(second) + self.current_indent) + ) + self.write(f"{lines[0]}\n") + + for line in lines[1:]: + self.write( + f"{'':>{second_col + self.current_indent * 4}}{line}\n" + ) + else: + self.write("\n") + else: + raise TypeError( + "Expected either three or two columns for definition list" + ) From 9b2b3bf02d4de7eb4b75cfc247d161cf3fde5152 Mon Sep 17 00:00:00 2001 From: Alex Strick van Linschoten Date: Wed, 22 Oct 2025 10:35:52 +0200 Subject: [PATCH 24/25] Restore type ignore comment for write_dl call The base branch had a type: ignore[arg-type] comment because we pass 3-column tuples to write_dl which expects 2-column tuples. The old formatter implementation handled this internally. Since we reverted the formatter refactoring, we need to restore this type ignore comment. --- src/zenml/cli/cli.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/zenml/cli/cli.py b/src/zenml/cli/cli.py index 1e1a5c56f29..7eb99cee703 100644 --- a/src/zenml/cli/cli.py +++ b/src/zenml/cli/cli.py @@ -141,7 +141,7 @@ def format_commands( "[dim cyan]Available ZenML Commands (grouped)[/dim cyan]" ) with formatter.section(section_title): - formatter.write_dl(rows) + formatter.write_dl(rows) # type: ignore[arg-type] @click.group(cls=ZenMLCLI, invoke_without_command=True) From f89bb1431a5fb3d929ae4773e5682f5be1b54b95 Mon Sep 17 00:00:00 2001 From: Alex Strick van Linschoten Date: Thu, 23 Oct 2025 07:08:31 +0200 Subject: [PATCH 25/25] Fix test_reusing_private_secret_name_succeeds on REST stores The test was failing on docker-server-mysql configurations because it reused a store object captured before entering a UserContext. In REST store setups, the store reads credentials from GlobalConfiguration on every API call. When the old store variable was used inside the UserContext, it inadvertently authenticated as User 2 instead of User 1. The fix creates a fresh Client/store after exiting the UserContext, ensuring it picks up the restored User 1 credentials from the cleaned-up GlobalConfiguration. This allows the test to correctly verify that User 1 can still see their own private secret after User 2's context exits. --- .../zen_stores/test_secrets_store.py | 20 +++++++++++-------- 1 file changed, 12 insertions(+), 8 deletions(-) diff --git a/tests/integration/functional/zen_stores/test_secrets_store.py b/tests/integration/functional/zen_stores/test_secrets_store.py index 877622c8dd9..3f8c0970c82 100644 --- a/tests/integration/functional/zen_stores/test_secrets_store.py +++ b/tests/integration/functional/zen_stores/test_secrets_store.py @@ -557,14 +557,18 @@ def test_reusing_private_secret_name_succeeds(): assert len(private_secrets) == 1 assert other_secret.id == private_secrets[0].id - other_private_secrets = store.list_secrets( - SecretFilter( - name=secret.name, - private=True, - ), - ).items - assert len(other_private_secrets) == 1 - assert secret.id == other_private_secrets[0].id + # After switching back to the original user context, instantiate a fresh client/store + original_client = Client() + original_store = original_client.zen_store + + other_private_secrets = original_store.list_secrets( + SecretFilter( + name=secret.name, + private=True, + ), + ).items + assert len(other_private_secrets) == 1 + assert secret.id == other_private_secrets[0].id def test_update_scope_succeeds():