|
"""Custom Help Formatters for Anomalib CLI.""" |
|
|
|
|
|
|
|
|
|
import argparse |
|
import re |
|
import sys |
|
from typing import TypeVar |
|
|
|
import docstring_parser |
|
from jsonargparse import DefaultHelpFormatter |
|
from rich.markdown import Markdown |
|
from rich.panel import Panel |
|
from rich_argparse import RichHelpFormatter |
|
|
|
REQUIRED_ARGUMENTS = { |
|
"train": {"model", "model.help", "data", "data.help", "ckpt_path", "config"}, |
|
"fit": {"model", "model.help", "data", "data.help", "ckpt_path", "config"}, |
|
"validate": {"model", "model.help", "data", "data.help", "ckpt_path", "config"}, |
|
"test": {"model", "model.help", "data", "data.help", "ckpt_path", "config"}, |
|
"predict": {"model", "model.help", "data", "data.help", "ckpt_path", "config"}, |
|
"export": {"model", "model.help", "export_type", "ckpt_path", "config"}, |
|
} |
|
|
|
try: |
|
from anomalib.engine import Engine |
|
|
|
DOCSTRING_USAGE = { |
|
"train": Engine.train, |
|
"fit": Engine.fit, |
|
"validate": Engine.validate, |
|
"test": Engine.test, |
|
"predict": Engine.predict, |
|
"export": Engine.export, |
|
} |
|
except ImportError: |
|
print("To use other subcommand using `anomalib install`") |
|
|
|
|
|
def get_short_docstring(component: TypeVar) -> str: |
|
"""Get the short description from the docstring. |
|
|
|
Args: |
|
component (TypeVar): The component to get the docstring from |
|
|
|
Returns: |
|
str: The short description |
|
""" |
|
if component.__doc__ is None: |
|
return "" |
|
docstring = docstring_parser.parse(component.__doc__) |
|
return docstring.short_description |
|
|
|
|
|
def get_verbosity_subcommand() -> dict: |
|
"""Parse command line arguments and returns a dictionary of key-value pairs. |
|
|
|
Returns: |
|
A dictionary containing the parsed command line arguments. |
|
|
|
Examples: |
|
>>> import sys |
|
>>> sys.argv = ['anomalib', 'train', '-h', '-v'] |
|
>>> get_verbosity_subcommand() |
|
{'subcommand': 'train', 'help': True, 'verbosity': 1} |
|
""" |
|
arguments: dict = {"subcommand": None, "help": False, "verbosity": 2} |
|
if len(sys.argv) >= 2 and sys.argv[1] not in ("--help", "-h"): |
|
arguments["subcommand"] = sys.argv[1] |
|
if "--help" in sys.argv or "-h" in sys.argv: |
|
arguments["help"] = True |
|
if arguments["subcommand"] in REQUIRED_ARGUMENTS: |
|
arguments["verbosity"] = 0 |
|
if "-v" in sys.argv or "--verbose" in sys.argv: |
|
arguments["verbosity"] = 1 |
|
if "-vv" in sys.argv: |
|
arguments["verbosity"] = 2 |
|
return arguments |
|
|
|
|
|
def get_intro() -> Markdown: |
|
"""Return a Markdown object containing the introduction text for Anomalib CLI Guide. |
|
|
|
The introduction text includes a brief description of the guide and links to the Github repository and documentation |
|
|
|
Returns: |
|
A Markdown object containing the introduction text for Anomalib CLI Guide. |
|
""" |
|
intro_markdown = ( |
|
"# Anomalib CLI Guide\n\n" |
|
"Github Repository: [https://github.com/openvinotoolkit/anomalib](https://github.com/openvinotoolkit/anomalib)." |
|
"\n\n" |
|
"A better guide is provided by the [documentation](https://anomalib.readthedocs.io/en/latest/index.html)." |
|
) |
|
return Markdown(intro_markdown) |
|
|
|
|
|
def get_verbose_usage(subcommand: str = "train") -> str: |
|
"""Return a string containing verbose usage information for the specified subcommand. |
|
|
|
Args: |
|
---- |
|
subcommand (str): The name of the subcommand to get verbose usage information for. Defaults to "train". |
|
|
|
Returns: |
|
------- |
|
str: A string containing verbose usage information for the specified subcommand. |
|
""" |
|
return ( |
|
"To get more overridable argument information, run the command below.\n" |
|
"```python\n" |
|
"# Verbosity Level 1\n" |
|
f"anomalib {subcommand} [optional_arguments] -h -v\n" |
|
"# Verbosity Level 2\n" |
|
f"anomalib {subcommand} [optional_arguments] -h -vv\n" |
|
"```" |
|
) |
|
|
|
|
|
def get_cli_usage_docstring(component: object | None) -> str | None: |
|
r"""Get the cli usage from the docstring. |
|
|
|
Args: |
|
---- |
|
component (Optional[object]): The component to get the docstring from |
|
|
|
Returns: |
|
------- |
|
Optional[str]: The quick-start guide as Markdown format. |
|
|
|
Example: |
|
------- |
|
component.__doc__ = ''' |
|
<Prev Section> |
|
|
|
CLI Usage: |
|
1. First Step. |
|
2. Second Step. |
|
|
|
<Next Section> |
|
''' |
|
>>> get_cli_usage_docstring(component) |
|
"1. First Step.\n2. Second Step." |
|
""" |
|
if component is None or component.__doc__ is None or "CLI Usage" not in component.__doc__: |
|
return None |
|
|
|
pattern = r"CLI Usage:(.*?)(?=\n{2,}|\Z)" |
|
match = re.search(pattern, component.__doc__, re.DOTALL) |
|
|
|
if match: |
|
contents = match.group(1).strip().split("\n") |
|
return "\n".join([content.strip() for content in contents]) |
|
return None |
|
|
|
|
|
def render_guide(subcommand: str | None = None) -> list: |
|
"""Render a guide for the specified subcommand. |
|
|
|
Args: |
|
---- |
|
subcommand (Optional[str]): The subcommand to render the guide for. |
|
|
|
Returns: |
|
------- |
|
list: A list of contents to be displayed in the guide. |
|
""" |
|
if subcommand is None or subcommand not in DOCSTRING_USAGE: |
|
return [] |
|
contents = [get_intro()] |
|
target_command = DOCSTRING_USAGE[subcommand] |
|
cli_usage = get_cli_usage_docstring(target_command) |
|
if cli_usage is not None: |
|
cli_usage += f"\n{get_verbose_usage(subcommand)}" |
|
quick_start = Panel(Markdown(cli_usage), border_style="dim", title="Quick-Start", title_align="left") |
|
contents.append(quick_start) |
|
return contents |
|
|
|
|
|
class CustomHelpFormatter(RichHelpFormatter, DefaultHelpFormatter): |
|
"""A custom help formatter for Anomalib CLI. |
|
|
|
This formatter extends the RichHelpFormatter and DefaultHelpFormatter classes to provide |
|
a more detailed and customizable help output for Anomalib CLI. |
|
|
|
Attributes: |
|
verbosity_level : int |
|
The level of verbosity for the help output. |
|
subcommand : str | None |
|
The subcommand to render the guide for. |
|
|
|
Methods: |
|
add_usage(usage, actions, *args, **kwargs) |
|
Add usage information to the help output. |
|
add_argument(action) |
|
Add an argument to the help output. |
|
format_help() |
|
Format the help output. |
|
""" |
|
|
|
verbosity_dict = get_verbosity_subcommand() |
|
verbosity_level = verbosity_dict["verbosity"] |
|
subcommand = verbosity_dict["subcommand"] |
|
|
|
def add_usage(self, usage: str | None, actions: list, *args, **kwargs) -> None: |
|
"""Add usage information to the formatter. |
|
|
|
Args: |
|
---- |
|
usage (str | None): A string describing the usage of the program. |
|
actions (list): An list of argparse.Action objects. |
|
*args (Any): Additional positional arguments to pass to the superclass method. |
|
**kwargs (Any): Additional keyword arguments to pass to the superclass method. |
|
|
|
Returns: |
|
------- |
|
None |
|
""" |
|
if self.subcommand in REQUIRED_ARGUMENTS: |
|
if self.verbosity_level == 0: |
|
actions = [] |
|
elif self.verbosity_level == 1: |
|
actions = [action for action in actions if action.dest in REQUIRED_ARGUMENTS[self.subcommand]] |
|
|
|
super().add_usage(usage, actions, *args, **kwargs) |
|
|
|
def add_argument(self, action: argparse.Action) -> None: |
|
"""Add an argument to the help formatter. |
|
|
|
If the verbose level is set to 0, the argument is not added. |
|
If the verbose level is set to 1 and the argument is not in the non-skip list, the argument is not added. |
|
|
|
Args: |
|
---- |
|
action (argparse.Action): The action to add to the help formatter. |
|
""" |
|
if self.subcommand in REQUIRED_ARGUMENTS: |
|
if self.verbosity_level == 0: |
|
return |
|
if self.verbosity_level == 1 and action.dest not in REQUIRED_ARGUMENTS[self.subcommand]: |
|
return |
|
super().add_argument(action) |
|
|
|
def format_help(self) -> str: |
|
"""Format the help message for the current command and returns it as a string. |
|
|
|
The help message includes information about the command's arguments and options, |
|
as well as any additional information provided by the command's help guide. |
|
|
|
Returns: |
|
str: A string containing the formatted help message. |
|
""" |
|
with self.console.capture() as capture: |
|
section = self._root_section |
|
if self.subcommand in REQUIRED_ARGUMENTS and self.verbosity_level in (0, 1) and len(section.rich_items) > 1: |
|
contents = render_guide(self.subcommand) |
|
for content in contents: |
|
self.console.print(content) |
|
if self.verbosity_level > 0: |
|
if len(section.rich_items) > 1: |
|
section = Panel(section, border_style="dim", title="Arguments", title_align="left") |
|
self.console.print(section, highlight=False, soft_wrap=True) |
|
help_msg = capture.get() |
|
|
|
if help_msg: |
|
help_msg = self._long_break_matcher.sub("\n\n", help_msg).rstrip() + "\n" |
|
return help_msg |
|
|