diff options
author | Daniel Baumann <daniel.baumann@progress-linux.org> | 2024-03-28 06:11:39 +0000 |
---|---|---|
committer | Daniel Baumann <daniel.baumann@progress-linux.org> | 2024-03-28 06:11:39 +0000 |
commit | 1fd6a618b60d7168fd8f37585d5d39d22d775afd (patch) | |
tree | fbc6d0c213b8acdd0a31deafe5c5fc0d05a3a312 /docs/scripts/generate_svg.py | |
parent | Initial commit. (diff) | |
download | anta-1fd6a618b60d7168fd8f37585d5d39d22d775afd.tar.xz anta-1fd6a618b60d7168fd8f37585d5d39d22d775afd.zip |
Adding upstream version 0.13.0.upstream/0.13.0
Signed-off-by: Daniel Baumann <daniel.baumann@progress-linux.org>
Diffstat (limited to 'docs/scripts/generate_svg.py')
-rw-r--r-- | docs/scripts/generate_svg.py | 92 |
1 files changed, 92 insertions, 0 deletions
diff --git a/docs/scripts/generate_svg.py b/docs/scripts/generate_svg.py new file mode 100644 index 0000000..19177db --- /dev/null +++ b/docs/scripts/generate_svg.py @@ -0,0 +1,92 @@ +# Copyright (c) 2023-2024 Arista Networks, Inc. +# Use of this source code is governed by the Apache License 2.0 +# that can be found in the LICENSE file. +""" +A script to generate svg files from anta command + +usage: + +python generate_svg.py anta ... +""" + +import io +import os +import pathlib +import sys +from contextlib import redirect_stdout, suppress +from importlib import import_module +from importlib.metadata import entry_points +from unittest.mock import patch + +from rich.console import Console + +from anta.cli.console import console +from anta.cli.nrfu.utils import anta_progress_bar + +OUTPUT_DIR = pathlib.Path(__file__).parent.parent / "imgs" + + +def custom_progress_bar() -> None: + """ + Set the console of progress_bar to main anta console + + Caveat: this capture all steps of the progress bar.. + Disabling refresh to only capture beginning and end + """ + progress = anta_progress_bar() + progress.live.auto_refresh = False + progress.live.console = console + return progress + + +if __name__ == "__main__": + # Sane rich size + os.environ["COLUMNS"] = "165" + + # stolen from https://github.com/ewels/rich-click/blob/main/src/rich_click/cli.py + args = sys.argv[1:] + script_name = args[0] + scripts = {script.name: script for script in entry_points().get("console_scripts")} + + if script_name in scripts: + # A VALID SCRIPT WAS passed + script = scripts[script_name] + module_path, function_name = script.value.split(":", 1) + prog = script_name + elif ":" in script_name: + # the path to a function was passed + module_path, function_name = args[0].split(":", 1) + prog = module_path.split(".", 1)[0] + else: + print("This is supposed to be used with anta only") + print("Usage: python generate_svg.py anta <options>") + sys.exit(1) + + sys.argv = [prog, *args[1:]] + module = import_module(module_path) + function = getattr(module, function_name) + + # Console to captur everything + new_console = Console(record=True) + + # tweaks to record and redirect to a dummy file + pipe = io.StringIO() + console.record = True + console.file = pipe + + # Redirect stdout of the program towards another StringIO to capture help + # that is not part or anta rich console + with redirect_stdout(io.StringIO()) as f: + # redirect potential progress bar output to console by patching + with patch("anta.cli.nrfu.commands.anta_progress_bar", custom_progress_bar): + with suppress(SystemExit): + function() + # print to our new console the output of anta console + new_console.print(console.export_text()) + # print the content of the stdout to our new_console + new_console.print(f.getvalue()) + + filename = f"{'_'.join(map(lambda x: x.replace('/', '_').replace('-', '_').replace('.', '_'), args))}.svg" + filename = f"{OUTPUT_DIR}/{filename}" + print(f"File saved at {filename}") + new_console.save_svg(filename, title=" ".join(args)) |