feat: Merge Recorder and OutputManager. (#182)

* feat: Merge Recorder and OutputManager. This PR seeks to merge Recorder and OutputManager into a unified interface for processing commands and producing output, either to file or console. This ensures that: - Filtering (`|`or `|!`) works from commands called via `source` - Output to files is filtered as expected. - Removes recorder.py and its class. - Incorporates recording into OutputManager. - The OutputManager handles recording (mostly) behind the scenes if we're in recording mode. Outside of log handling we don't have to call recording tools explicitly as its done as part of normal output processing. - Incorporates table formatting into OutputManager where it belongs. - Moves command processing to a shared separate function called from both `source` and the interactive loop. - Adds commands to start and stop recording interactively. - Updates the prompt to show active recording if applicable. NOTE: This PR does not pass the current tests as it fixes #181. These changes should be thourughly reviewed and verified before new test logs are created. Fixes #179, fixes #181. * Add back a comment about shell commands called from scripts. * Clarify comments about recording related attributes. * Avoid reimplementing stop_recording when saving. * Update testsuite results.
unioslo · Dec 6, 2023 · c37d919 · c37d919
1 parent a7d33bc
commit c37d919
Show file tree

Hide file tree

Showing 11 changed files with 528 additions and 270 deletions.
diff --git a/ci/testsuite-result.json b/ci/testsuite-result.json
diff --git a/mreg_cli/bacnet.py b/mreg_cli/bacnet.py
@@ -2,7 +2,8 @@
 from .history import history
 from .host import host
 from .log import cli_error, cli_info
-from .util import add_formatted_table_for_output, delete, get, get_list, host_info_by_name, post
+from .outputmanager import OutputManager
+from .util import delete, get, get_list, host_info_by_name, post
 
 
 def bacnetid_add(args) -> None:
@@ -74,7 +75,7 @@ def bacnetid_list(args) -> None:
         if maxval > 4194302:
             cli_error("The maximum ID value is 4194302.")
     r = get_list("/api/v1/bacnet/ids/", {"id__range": "{},{}".format(minval, maxval)})
-    add_formatted_table_for_output(("ID", "Hostname"), ("id", "hostname"), r)
+    OutputManager().add_formatted_table(("ID", "Hostname"), ("id", "hostname"), r)
 
 
 host.add_command(

diff --git a/mreg_cli/cli.py b/mreg_cli/cli.py
@@ -2,12 +2,15 @@
 import html
 import os
 import shlex
+import sys
+from typing import Generator, List
 
 from prompt_toolkit import HTML, print_formatted_text
 from prompt_toolkit.completion import Completer, Completion
 
-from . import recorder, util
+from . import util
 from .exceptions import CliError, CliWarning
+from .outputmanager import OutputManager, remove_comments
 
 
 class CliExit(Exception):
@@ -143,9 +146,9 @@ def parse(self, command: str) -> None:
             exc.print_self()
 
         except CliExit:
-            from sys import exit
-
-            exit(0)
+            # If we have active recordings going on, save them before exiting
+            OutputManager().save_recording()
+            sys.exit(0)
 
         else:
             # If no exception occurred make sure errno isn't set to an error
@@ -207,6 +210,21 @@ def complete(self, cur, words):
                         start_position=-len(cur),
                     )
 
+    def process_command_line(self, line: str) -> None:
+        """Process a line containing a command."""
+        line = remove_comments(line)
+        # OutputManager is a singleton class so we
+        # need to clear it before each command.
+        output = OutputManager()
+        output.clear()
+        # Set the command that generated the output
+        # Also remove filters and other noise.
+        cmd = output.from_command(line)
+        # Run the command
+        cli.parse(cmd)
+        # Render the output
+        output.render()
+
 
 # Top parser is the root of all the command parsers
 _top_parser = argparse.ArgumentParser("")
@@ -217,6 +235,19 @@ def _quit(args):
     raise CliExit
 
 
+def _start_recording(args) -> None:
+    """Start recording commands and output to the given file."""
+    if not args.filename:
+        raise CliError("No filename given.")
+
+    OutputManager().start_recording(args.filename)
+
+
+def _stop_recording(args):
+    """Stop recording commands and output to the given file."""
+    OutputManager().save_recording()
+
+
 # Always need a quit command
 cli.add_command(
     prog="quit",
@@ -245,16 +276,44 @@ def logout(args):
     callback=logout,
 )
 
+recordings = cli.add_command(
+    prog="recording",
+    description="Recording related commands.",
+    short_desc="Recording related commands",
+)
 
-def source(files, ignore_errors, verbose):
-    """Source reads commands from one or more source files.
-    Each command must be on one line and the commands must be separated with
-    newlines.
-    The files may contain comments. The comment symbol is #.
-    """
+recordings.add_command(
+    prog="start",
+    description="Start recording commands to a file.",
+    short_desc="Start recording",
+    callback=_start_recording,
+    flags=[
+        Flag(
+            "filename",
+            description="The filename to record to.",
+            short_desc="Filename",
+            metavar="filename",
+        )
+    ],
+)
+
+recordings.add_command(
+    prog="stop",
+    description="Stop recording commands and output to the given file.",
+    short_desc="Stop recording",
+    callback=_stop_recording,
+)
 
-    rec = recorder.Recorder()
 
+def source(files: List[str], ignore_errors: bool, verbose: bool) -> Generator[str, None, None]:
+    """Read commands from one or more source files and yield them.
+
+    :param files: List of file paths to read commands from.
+    :param ignore_errors: If True, continue on errors.
+    :param verbose: If True, print commands before execution.
+
+    :yields: Command lines from the files.
+    """
     for filename in files:
         if filename.startswith("~"):
             filename = os.path.expanduser(filename)
@@ -266,15 +325,11 @@ def source(files, ignore_errors, verbose):
                         os.system(line[1:])
                         continue
 
-                    # If recording commands, submit the command line.
-                    # Don't record the "source" command itself.
-                    if rec.is_recording() and not line.lstrip().startswith("source"):
-                        rec.record_command(line)
-
-                    # In verbose mode all commands are printed before execution.
                     if verbose:
                         print_formatted_text(HTML(f"<i>> {html.escape(line.strip())}</i>"))
-                    cli.parse(line)
+
+                    yield line
+
                     if cli.last_errno != 0:
                         print_formatted_text(
                             HTML(
@@ -284,6 +339,7 @@ def source(files, ignore_errors, verbose):
                                 )
                             )
                         )
+                        OutputManager().record_extra_output(f"{filename}: Error on line {i + 1}")
                         if not ignore_errors:
                             return
         except FileNotFoundError:
@@ -293,7 +349,14 @@ def source(files, ignore_errors, verbose):
 
 
 def _source(args):
-    source(args.files, args.ignore_errors, args.verbose)
+    """Wrapper for the source function to integrate with the CLI.
+
+    :param args: Arguments from the CLI.
+    """
+    for command in source(args.files, args.ignore_errors, args.verbose):
+        # Process each command here as needed, similar to the main loop
+        print(f"Processing command: {command}")
+        cli.process_command_line(command)
 
 
 # Always need the source command.

diff --git a/mreg_cli/label.py b/mreg_cli/label.py
@@ -2,7 +2,7 @@
 from .history import history
 from .log import cli_info, cli_warning
 from .outputmanager import OutputManager
-from .util import add_formatted_table_for_output, delete, get, get_list, patch, post
+from .util import delete, get, get_list, patch, post
 
 label = cli.add_command(
     prog="label",
@@ -38,7 +38,7 @@ def label_list(args) -> None:
     if not labels:
         cli_info("No labels", True)
         return
-    add_formatted_table_for_output(("Name", "Description"), ("name", "description"), labels)
+    OutputManager().add_formatted_table(("Name", "Description"), ("name", "description"), labels)
 
 
 label.add_command(prog="list", description="List labels", callback=label_list, flags=[])
@@ -83,7 +83,7 @@ def label_info(args) -> None:
     permlist = get_list("/api/v1/permissions/netgroupregex/", params={"labels__name": args.name})
     manager.add_line("Permissions with this label:")
     if permlist:
-        add_formatted_table_for_output(
+        OutputManager().add_formatted_table(
             ("IP range", "Group", "Reg.exp."),
             ("range", "group", "regex"),
             permlist,

diff --git a/mreg_cli/log.py b/mreg_cli/log.py
@@ -4,8 +4,8 @@
 from datetime import datetime
 from typing import NoReturn, Optional, Type
 
-from . import recorder
 from .exceptions import CliError, CliWarning
+from .outputmanager import OutputManager
 
 logfile = None
 
@@ -45,10 +45,10 @@ def cli_error(
     if raise_exception:
         # A simplified message for console
         msg = "ERROR: {}: {}".format(pre, msg)
-        rec = recorder.Recorder()
-        if rec.is_recording():
+        manager = OutputManager()
+        if manager.is_recording():
             # If recording traffic, also record the console output
-            rec.record_output(msg)
+            manager.record_extra_output(msg)
         # Raise the exception
         raise exception(msg)
     return None
@@ -69,10 +69,10 @@ def cli_warning(
     if raise_exception:
         # A simplified message for console
         msg = "WARNING: {}: {}".format(pre, msg)
-        rec = recorder.Recorder()
-        if rec.is_recording():
+        manager = OutputManager()
+        if manager.is_recording():
             # If recording traffic, also record the console output
-            rec.record_output(msg)
+            manager.record_extra_output(msg)
         raise exception(msg)
     return None
 
@@ -91,7 +91,7 @@ def cli_info(msg: str, print_msg: bool = False) -> None:
         # A simplified message for console
         msg = "OK: {}: {}".format(pre, msg)
         print(msg)
-        rec = recorder.Recorder()
-        if rec.is_recording():
+        manager = OutputManager()
+        if manager.is_recording():
             # If recording traffic, also record the console output
-            rec.record_output(msg)
+            manager.record_extra_output(msg)
diff --git a/mreg_cli/main.py b/mreg_cli/main.py
@@ -6,7 +6,7 @@
 from prompt_toolkit import HTML
 from prompt_toolkit.shortcuts import CompleteStyle, PromptSession
 
-from . import config, log, recorder, util
+from . import config, log, util
 from .cli import cli, source
 from .exceptions import CliError, CliWarning
 from .outputmanager import OutputManager
@@ -112,9 +112,8 @@ def main():
     if "logfile" in conf:
         log.logfile = conf["logfile"]
 
-    rec = recorder.Recorder()
     if "record_traffic" in conf:
-        rec.start_recording(conf["record_traffic"])
+        OutputManager().start_recording(conf["record_traffic"])
 
     if "user" not in conf:
         print("Username not set in config or as argument")
@@ -144,11 +143,20 @@ def main():
     from . import policy  # noqa: F401
     from . import zone  # noqa: F401
 
+    # Define a function that returns the prompt message
+    def get_prompt_message():
+        """Return the prompt message."""
+        manager = OutputManager()
+        if manager.is_recording():
+            return HTML(f"<i>[>'{manager.recording_filename()}']</i> <b>{args.prompt}</b>> ")
+        else:
+            return HTML(f"<b>{args.prompt}</b>> ")
+
     # session is a PromptSession object from prompt_toolkit which handles
     # some configurations of the prompt for us: the text of the prompt; the
     # completer; and other visual things.
     session = PromptSession(
-        message=HTML(f"<b>{args.prompt}</b>> "),
+        message=get_prompt_message,
         search_ignore_case=True,
         completer=cli,
         complete_while_typing=True,
@@ -160,7 +168,8 @@ def main():
 
     # if the --source parameter was given, read commands from the source file and then exit
     if "source" in conf:
-        source([conf["source"]], "verbosity" in conf, False)
+        for command in source([conf["source"]], "verbosity" in conf, False):
+            cli.process_command_line(command)
         return
 
     # The app runs in an infinite loop and is expected to exit using sys.exit()
@@ -173,25 +182,9 @@ def main():
             raise SystemExit() from None
         try:
             for line in lines.splitlines():
-                # If recording commands, submit the command line.
-                # Don't record the "source" command itself.
-                if rec.is_recording() and not line.lstrip().startswith("source"):
-                    rec.record_command(line)
-                # OutputManager is a singleton class so we
-                # need to clear it before each command.
-                output = OutputManager()
-                output.clear()
-                # Set the command that generated the output
-                # Also remove filters and other noise.
-                cmd = output.from_command(line)
-                # Run the command
-                cli.parse(cmd)
-                # Render the output
-                output.render()
-        except (CliWarning, CliError) as exc:
-            exc.print_self()
-        except ValueError as exc:
-            print(exc)
+                cli.process_command_line(line)
+        except ValueError as e:
+            print(e)
 
 
 if __name__ == "__main__":