Skip to content

Cli

zenml.cli.cli

Core CLI functionality.

TagGroup (Group)

Override the default click Group to add a tag.

The tag is used to group commands and groups of commands in the help output.

Source code in zenml/cli/cli.py
class TagGroup(click.Group):
    """Override the default click Group to add a tag.

    The tag is used to group commands and groups of
    commands in the help output.
    """

    def __init__(
        self,
        name: Optional[str] = None,
        tag: Optional[CliCategories] = None,
        commands: Optional[
            Union[Dict[str, click.Command], Sequence[click.Command]]
        ] = None,
        **kwargs: Dict[str, Any],
    ) -> None:
        """Initialize the Tag group.

        Args:
            name: The name of the group.
            tag: The tag of the group.
            commands: The commands of the group.
            kwargs: Additional keyword arguments.
        """
        super(TagGroup, self).__init__(name, commands, **kwargs)
        self.tag = tag or CliCategories.OTHER_COMMANDS

__init__(self, name=None, tag=None, commands=None, **kwargs) special

Initialize the Tag group.

Parameters:

Name Type Description Default
name Optional[str]

The name of the group.

None
tag Optional[zenml.enums.CliCategories]

The tag of the group.

None
commands Union[Dict[str, click.core.Command], Sequence[click.core.Command]]

The commands of the group.

None
kwargs Dict[str, Any]

Additional keyword arguments.

{}
Source code in zenml/cli/cli.py
def __init__(
    self,
    name: Optional[str] = None,
    tag: Optional[CliCategories] = None,
    commands: Optional[
        Union[Dict[str, click.Command], Sequence[click.Command]]
    ] = None,
    **kwargs: Dict[str, Any],
) -> None:
    """Initialize the Tag group.

    Args:
        name: The name of the group.
        tag: The tag of the group.
        commands: The commands of the group.
        kwargs: Additional keyword arguments.
    """
    super(TagGroup, self).__init__(name, commands, **kwargs)
    self.tag = tag or CliCategories.OTHER_COMMANDS

ZenContext (Context)

Override the default click Context to add the new Formatter.

Source code in zenml/cli/cli.py
class ZenContext(click.Context):
    """Override the default click Context to add the new Formatter."""

    formatter_class = ZenFormatter

formatter_class (HelpFormatter)

Override the default HelpFormatter to add a custom format for the help command output.

Source code in zenml/cli/cli.py
class ZenFormatter(formatting.HelpFormatter):
    """Override the default HelpFormatter to add a custom format for the help command output."""

    def __init__(
        self,
        indent_increment: int = 2,
        width: Optional[int] = None,
        max_width: Optional[int] = None,
    ) -> None:
        """Initialize the formatter.

        Args:
            indent_increment: The number of spaces to indent each level of
                nesting.
            width: The maximum width of the help output.
            max_width: The maximum width of the help output.
        """
        super(ZenFormatter, self).__init__(indent_increment, width, max_width)
        self.current_indent = 0

    def write_dl(
        self,
        rows: Sequence[Tuple[str, ...]],
        col_max: int = 30,
        col_spacing: int = 2,
    ) -> None:
        """Writes a definition list into the buffer.

        This is how options and commands are usually formatted.

        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 the format
                            ``(term, value, description)``.

        Raises:
            TypeError: if the number of columns is not 2 or 3.
        """
        rows = list(rows)
        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
                )
                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
                )
                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"
            )
__init__(self, indent_increment=2, width=None, max_width=None) special

Initialize the formatter.

Parameters:

Name Type Description Default
indent_increment int

The number of spaces to indent each level of nesting.

2
width Optional[int]

The maximum width of the help output.

None
max_width Optional[int]

The maximum width of the help output.

None
Source code in zenml/cli/cli.py
def __init__(
    self,
    indent_increment: int = 2,
    width: Optional[int] = None,
    max_width: Optional[int] = None,
) -> None:
    """Initialize the formatter.

    Args:
        indent_increment: The number of spaces to indent each level of
            nesting.
        width: The maximum width of the help output.
        max_width: The maximum width of the help output.
    """
    super(ZenFormatter, self).__init__(indent_increment, width, max_width)
    self.current_indent = 0
write_dl(self, rows, col_max=30, col_spacing=2)

Writes a definition list into the buffer.

This is how options and commands are usually formatted.

Parameters:

Name Type Description Default
rows Sequence[Tuple[str, ...]]

a list of items as tuples for the terms and values.

required
col_max int

the maximum width of the first column.

30
col_spacing int

the number of spaces between the first and second column (and third).

2

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 the format (term, value, description).

Exceptions:

Type Description
TypeError

if the number of columns is not 2 or 3.

Source code in zenml/cli/cli.py
def write_dl(
    self,
    rows: Sequence[Tuple[str, ...]],
    col_max: int = 30,
    col_spacing: int = 2,
) -> None:
    """Writes a definition list into the buffer.

    This is how options and commands are usually formatted.

    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 the format
                        ``(term, value, description)``.

    Raises:
        TypeError: if the number of columns is not 2 or 3.
    """
    rows = list(rows)
    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
            )
            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
            )
            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"
        )

ZenMLCLI (Group)

Override the default click Group to create a custom format command help output.

Source code in zenml/cli/cli.py
class ZenMLCLI(click.Group):
    """Override the default click Group to create a custom format command help output."""

    context_class = ZenContext

    def get_help(self, ctx: Context) -> str:
        """Formats the help into a string and returns it.

        Calls :meth:`format_help` internally.

        Args:
            ctx: The click context.

        Returns:
            The formatted help string.
        """
        formatter = ctx.make_formatter()
        self.format_help(ctx, formatter)
        # TODO [ENG-862]: Find solution for support console.pager and color support in print
        rich.print(formatter.getvalue().rstrip("\n"))
        return ""

    def format_commands(
        self, ctx: click.Context, formatter: formatting.HelpFormatter
    ) -> None:
        """Extra format methods for multi methods that adds all the commands after the options.

        This custom format_commands method is used to retrieve the commands and
        groups of commands with a tag. In order to call the new custom format
        method, the command must be added to the ZenML CLI class.

        Args:
            ctx: The click context.
            formatter: The click formatter.
        """
        commands: List[Tuple[CliCategories, str, Union[Command, TagGroup]]] = []
        for subcommand in self.list_commands(ctx):
            cmd = self.get_command(ctx, subcommand)
            # What is this, the tool lied about a command.  Ignore it
            if cmd is None or cmd.hidden:
                continue
            category = (
                cmd.tag
                if isinstance(cmd, TagGroup)
                else CliCategories.OTHER_COMMANDS
            )
            commands.append(
                (
                    category,
                    subcommand,
                    cmd,
                )
            )

        if len(commands):
            ordered_categories = list(CliCategories.__members__.values())
            commands = list(
                sorted(
                    (commands),
                    key=lambda x: (
                        ordered_categories.index(x[0]),
                        x[0].value,
                        x[1],
                    ),
                )
            )
            rows: List[Tuple[str, str, str]] = []
            for (tag, subcommand, cmd) in 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]

context_class (Context)

Override the default click Context to add the new Formatter.

Source code in zenml/cli/cli.py
class ZenContext(click.Context):
    """Override the default click Context to add the new Formatter."""

    formatter_class = ZenFormatter
formatter_class (HelpFormatter)

Override the default HelpFormatter to add a custom format for the help command output.

Source code in zenml/cli/cli.py
class ZenFormatter(formatting.HelpFormatter):
    """Override the default HelpFormatter to add a custom format for the help command output."""

    def __init__(
        self,
        indent_increment: int = 2,
        width: Optional[int] = None,
        max_width: Optional[int] = None,
    ) -> None:
        """Initialize the formatter.

        Args:
            indent_increment: The number of spaces to indent each level of
                nesting.
            width: The maximum width of the help output.
            max_width: The maximum width of the help output.
        """
        super(ZenFormatter, self).__init__(indent_increment, width, max_width)
        self.current_indent = 0

    def write_dl(
        self,
        rows: Sequence[Tuple[str, ...]],
        col_max: int = 30,
        col_spacing: int = 2,
    ) -> None:
        """Writes a definition list into the buffer.

        This is how options and commands are usually formatted.

        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 the format
                            ``(term, value, description)``.

        Raises:
            TypeError: if the number of columns is not 2 or 3.
        """
        rows = list(rows)
        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
                )
                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
                )
                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"
            )
__init__(self, indent_increment=2, width=None, max_width=None) special

Initialize the formatter.

Parameters:

Name Type Description Default
indent_increment int

The number of spaces to indent each level of nesting.

2
width Optional[int]

The maximum width of the help output.

None
max_width Optional[int]

The maximum width of the help output.

None
Source code in zenml/cli/cli.py
def __init__(
    self,
    indent_increment: int = 2,
    width: Optional[int] = None,
    max_width: Optional[int] = None,
) -> None:
    """Initialize the formatter.

    Args:
        indent_increment: The number of spaces to indent each level of
            nesting.
        width: The maximum width of the help output.
        max_width: The maximum width of the help output.
    """
    super(ZenFormatter, self).__init__(indent_increment, width, max_width)
    self.current_indent = 0
write_dl(self, rows, col_max=30, col_spacing=2)

Writes a definition list into the buffer.

This is how options and commands are usually formatted.

Parameters:

Name Type Description Default
rows Sequence[Tuple[str, ...]]

a list of items as tuples for the terms and values.

required
col_max int

the maximum width of the first column.

30
col_spacing int

the number of spaces between the first and second column (and third).

2

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 the format (term, value, description).

Exceptions:

Type Description
TypeError

if the number of columns is not 2 or 3.

Source code in zenml/cli/cli.py
def write_dl(
    self,
    rows: Sequence[Tuple[str, ...]],
    col_max: int = 30,
    col_spacing: int = 2,
) -> None:
    """Writes a definition list into the buffer.

    This is how options and commands are usually formatted.

    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 the format
                        ``(term, value, description)``.

    Raises:
        TypeError: if the number of columns is not 2 or 3.
    """
    rows = list(rows)
    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
            )
            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
            )
            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"
        )

format_commands(self, ctx, formatter)

Extra format methods for multi methods that adds all the commands after the options.

This custom format_commands method is used to retrieve the commands and groups of commands with a tag. In order to call the new custom format method, the command must be added to the ZenML CLI class.

Parameters:

Name Type Description Default
ctx Context

The click context.

required
formatter HelpFormatter

The click formatter.

required
Source code in zenml/cli/cli.py
def format_commands(
    self, ctx: click.Context, formatter: formatting.HelpFormatter
) -> None:
    """Extra format methods for multi methods that adds all the commands after the options.

    This custom format_commands method is used to retrieve the commands and
    groups of commands with a tag. In order to call the new custom format
    method, the command must be added to the ZenML CLI class.

    Args:
        ctx: The click context.
        formatter: The click formatter.
    """
    commands: List[Tuple[CliCategories, str, Union[Command, TagGroup]]] = []
    for subcommand in self.list_commands(ctx):
        cmd = self.get_command(ctx, subcommand)
        # What is this, the tool lied about a command.  Ignore it
        if cmd is None or cmd.hidden:
            continue
        category = (
            cmd.tag
            if isinstance(cmd, TagGroup)
            else CliCategories.OTHER_COMMANDS
        )
        commands.append(
            (
                category,
                subcommand,
                cmd,
            )
        )

    if len(commands):
        ordered_categories = list(CliCategories.__members__.values())
        commands = list(
            sorted(
                (commands),
                key=lambda x: (
                    ordered_categories.index(x[0]),
                    x[0].value,
                    x[1],
                ),
            )
        )
        rows: List[Tuple[str, str, str]] = []
        for (tag, subcommand, cmd) in 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]

get_help(self, ctx)

Formats the help into a string and returns it.

Calls :meth:format_help internally.

Parameters:

Name Type Description Default
ctx Context

The click context.

required

Returns:

Type Description
str

The formatted help string.

Source code in zenml/cli/cli.py
def get_help(self, ctx: Context) -> str:
    """Formats the help into a string and returns it.

    Calls :meth:`format_help` internally.

    Args:
        ctx: The click context.

    Returns:
        The formatted help string.
    """
    formatter = ctx.make_formatter()
    self.format_help(ctx, formatter)
    # TODO [ENG-862]: Find solution for support console.pager and color support in print
    rich.print(formatter.getvalue().rstrip("\n"))
    return ""