Integrate presentation of the current state into ModuleHandler.

This partly revers commit 3e9d54d.
This reverts commit e15ce26.
This reverts commit bff688c.

Signed-off-by: Tobias Wolf <[email protected]>

Author: NotTheEvilOne

src/rookify/__main__.py

@@ -1,73 +1,15 @@
 # -*- coding: utf-8 -*-
 
-import json
 from pickle import Unpickler
 import sys
 import argparse
 from argparse import ArgumentParser
-from typing import Any, Dict, Optional
+from typing import Any, Dict
 from .modules import load_modules
 from .modules.machine import Machine
+from .modules.module import ModuleHandler
 from .logger import configure_logging, get_logger
 from .yaml import load_config
-from structlog.typing import BindableLogger
-from pathlib import Path
-
-
-def parse_args(args: list[str]) -> argparse.Namespace:
-    # Putting args-parser in seperate function to make this testable
-    arg_parser = ArgumentParser("Rookify")
-
-    # --dry-run option
-    arg_parser.add_argument("--dry-run", action="store_true", dest="dry_run_mode")
-
-    # --list-modules option
-    arg_parser.add_argument(
-        "--list-modules", action="store_true", help="List all modules"
-    )
-
-    # Custom ReadAction to set 'all' if nothing is specified for --read-pickle
-    class ReadAction(argparse.Action):
-        def __call__(
-            self,
-            parser: ArgumentParser,
-            namespace: argparse.Namespace,
-            values: Optional[Any],
-            option_string: Optional[str] = None,
-        ) -> None:
-            setattr(namespace, self.dest, values if values is not None else "all")
-
-    # Custom ShowProgressAction to set 'all' if nothing is specified for --show-progress
-    class ShowProgressAction(argparse.Action):
-        def __call__(
-            self,
-            parser: ArgumentParser,
-            namespace: argparse.Namespace,
-            values: Optional[Any],
-            option_string: Optional[str] = None,
-        ) -> None:
-            setattr(namespace, self.dest, values if values is not None else "all")
-
-    arg_parser.add_argument(
-        "--read-pickle",
-        nargs="?",
-        action=ReadAction,
-        dest="read_pickle",
-        metavar="<section>",
-        help="Show the content of the pickle file. Default argument is 'all', you can also specify a section you want to look at.",
-        required=False,
-    )
-
-    arg_parser.add_argument(
-        "--show-progress",
-        nargs="?",
-        action=ShowProgressAction,
-        dest="show_progress",
-        metavar="<module>",
-        help="Show progress of the modules. Default argument is 'all', you can also specify a module you want to get the progress status from.",
-        required=False,
-    )
-    return arg_parser.parse_args(args)
 
 
 def load_pickler(pickle_file_name: str) -> Any:
@@ -77,75 +19,26 @@ def load_pickler(pickle_file_name: str) -> Any:
         return states_data
 
 
-def get_all_modules() -> list[str]:
-    base_path = Path(__file__).resolve().parent
-    module_path = Path(base_path) / "modules"
-    module_names = []
-    for item in module_path.iterdir():
-        if item.is_dir() and item.name != "__pycache__":
-            module_names.append(item.name)
-    return module_names
-
-
-def sort_pickle_file(unsorted_states_data: Dict[str, Any]) -> Dict[str, Any]:
-    # sort the pickle-file alfabetically
-    iterable_dict = iter(unsorted_states_data)
-    first_key = next(iterable_dict)
-    data_values = unsorted_states_data[first_key]["data"]
-    sorted_data_by_keys = {k: data_values[k] for k in sorted(data_values)}
-    return sorted_data_by_keys
-
-
-def read_pickle_file(
-    args: argparse.Namespace, pickle_file_name: str, log: BindableLogger
-) -> None:
-    states_data = load_pickler(pickle_file_name)
-    sorted_states_data = sort_pickle_file(states_data)
+def parse_args(args: list[str]) -> argparse.Namespace:
+    # Putting args-parser in seperate function to make this testable
+    arg_parser = ArgumentParser("Rookify")
 
-    # Check if a specific section should be listed
-    if args.read_pickle != "all":
-        if args.read_pickle not in sorted_states_data.keys():
-            log.error(f"The section {args.read_pickle} does not exist")
-        else:
-            sorted_states_data = sorted_states_data[args.read_pickle]
+    # --dry-run option
+    arg_parser.add_argument("--dry-run", action="store_true", dest="dry_run_mode")
 
-    log.info(
-        'Current state as retrieved from pickle-file: \n "{0}": {1}'.format(
-            args.read_pickle, json.dumps(sorted_states_data, indent=4)
-        )
+    arg_parser.add_argument(
+        "--show-states",
+        action="store_true",
+        dest="show_states",
+        help="Show states of the modules.",
     )
 
-
-def show_progress_from_state(
-    args: argparse.Namespace, pickle_file_name: str, log: BindableLogger
-) -> bool:
-    # states_data = load_pickler(pickle_file_name)
-    modules = get_all_modules()
-
-    # Check if a specific module should be targeted
-    if args.show_progress != "all":
-        module = args.show_progress
-        if args.show_progress not in modules:
-            log.error(f"The module {module} does not exist")
-            return False
-        log.info("Show progress of the {0} module".format(args.show_progress))
-        return True
-    else:
-        log.info("Show progress of {0} modules".format(args.show_progress))
-        return True
+    return arg_parser.parse_args(args)
 
 
 def main() -> None:
     args = parse_args(sys.argv[1:])
 
-    # Handle --list-modules
-    if args.list_modules:
-        modules = get_all_modules()
-        print("Available modules:\n")
-        for module in modules:
-            print(f"- {module}")
-        return
-
     # Load configuration file
     try:
         config: Dict[str, Any] = load_config("config.yaml")
@@ -154,49 +47,29 @@ def main() -> None:
 
     # Configure logging
     try:
-        configure_logging(config["logging"])
+        if args.show_states is True:
+            configure_logging(
+                {"level": "ERROR", "format": {"renderer": "console", "time": "iso"}}
+            )
+        else:
+            configure_logging(config["logging"])
     except Exception as e:
         raise SystemExit(f"Error configuring logging: {e}")
+
     # Get Logger
     log = get_logger()
 
-    # Get Pickle File if configured in config.yaml
-    pickle_file_name = config["general"].get("machine_pickle_file")
-    if pickle_file_name is None:
-        log.info("No pickle file was set in the configuration.")
-    else:
-        log.info(f"Pickle file set: {pickle_file_name}")
+    log.info("Executing Rookify ...")
 
-    # Get Pickle File if configured in config.yaml
-    pickle_file_name = config["general"].get("machine_pickle_file")
-    if pickle_file_name is None:
-        log.info("No pickle file was set in the configuration.")
-    else:
-        log.info(f"Pickle file set: {pickle_file_name}")
-
-    # If read_pickle is run and there is a picklefile, show the picklefiles contents.
-    # NOTE: preflight mode (--dry-run) has no effect here, because no module actions are required.
-    if args.read_pickle is not None and pickle_file_name is not None:
-        read_pickle_file(args, pickle_file_name, log)
-        return
-    elif args.read_pickle is not None and pickle_file_name is None:
-        log.info(
-            "No pickle file configured to read from. Check if the pickle file exists and is configured in config.yaml"
-        )
-        return
+    machine = Machine(config["general"].get("machine_pickle_file"))
+
+    load_modules(machine, config)
 
+    if args.show_states is True:
+        ModuleHandler.show_states(machine, config)
     else:
-        machine = Machine(config["general"].get("machine_pickle_file"))
-
-        if args.show_progress is not None:
-            if show_progress_from_state(args, pickle_file_name, log) is True:
-                load_modules(machine, config, show_progress=True)
-                # NOTE: this is always run in preflight-mode (migration should not be executed)
-                machine.execute(dry_run_mode=True)
-                return
-            else:
-                return
-        else:
-            load_modules(machine, config)
-            log.debug("Executing Rookify")
-            machine.execute(dry_run_mode=args.dry_run_mode)
+        machine.execute(dry_run_mode=args.dry_run_mode)
+
+
+if __name__ == "__main__":
+    main()

src/rookify/modules/__init__.py

@@ -1,7 +1,7 @@
 # -*- coding: utf-8 -*-
 
 import importlib
-from typing import Any, Dict, Optional
+from typing import Any, Dict, List
 from ..logger import get_logger
 from .machine import Machine
 
@@ -22,11 +22,18 @@ def __init__(self, module_name: str, message: str):
         self.message = message
 
 
+_modules_loaded: List[Any] = []
+
+
+def get_modules() -> List[Any]:
+    global _modules_loaded
+    return _modules_loaded.copy()
+
+
 def _load_module(
     machine: Machine,
     config: Dict[str, Any],
     module_name: str,
-    show_progress: Optional[bool] = False,
 ) -> None:
     """
     Dynamically loads a module from the 'rookify.modules' package.
@@ -35,8 +42,10 @@ def _load_module(
     :return: returns tuple of preflight_modules, modules
     """
 
+    global _modules_loaded
+
     module = importlib.import_module("rookify.modules.{0}".format(module_name))
-    additional_modules = []
+    additional_module_names = []
 
     if not hasattr(module, "ModuleHandler") or not callable(
         getattr(module.ModuleHandler, "register_states")
@@ -45,17 +54,17 @@ def _load_module(
 
     if hasattr(module.ModuleHandler, "REQUIRES"):
         assert isinstance(module.ModuleHandler.REQUIRES, list)
-        additional_modules = module.ModuleHandler.REQUIRES
+        additional_module_names = module.ModuleHandler.REQUIRES
 
-    for module_name in additional_modules:
-        _load_module(machine, config, module_name, show_progress)
+    for additional_module_name in additional_module_names:
+        _load_module(machine, config, additional_module_name)
 
-    module.ModuleHandler.register_states(machine, config, show_progress)
+    if module not in _modules_loaded:
+        _modules_loaded.append(module)
+    module.ModuleHandler.register_states(machine, config)
 
 
-def load_modules(
-    machine: Machine, config: Dict[str, Any], show_progress: Optional[bool] = False
-) -> None:
+def load_modules(machine: Machine, config: Dict[str, Any]) -> None:
     """
     Dynamically loads modules from the 'modules' package.
 
@@ -68,7 +77,7 @@ def load_modules(
     for entry in importlib.resources.files("rookify.modules").iterdir():
         if entry.is_dir() and entry.name in config["migration_modules"]:
             migration_modules.remove(entry.name)
-            _load_module(machine, config, entry.name, show_progress)
+            _load_module(machine, config, entry.name)
 
     if len(migration_modules) > 0 or len(config["migration_modules"]) < 1:
         logger = get_logger()

src/rookify/modules/machine.py

@@ -26,22 +26,14 @@ def add_execution_state(self, name: str, **kwargs: Any) -> None:
     def add_preflight_state(self, name: str, **kwargs: Any) -> None:
         self._preflight_states.append(self.__class__.state_cls(name, **kwargs))
 
-    def execute(
-        self, dry_run_mode: bool = False, show_progress: Optional[bool] = False
-    ) -> None:
+    def execute(self, dry_run_mode: bool = False) -> None:
         states = self._preflight_states
         if not dry_run_mode:
             states = states + self._execution_states
 
-        logger = get_logger()
+        self._register_states(states)
 
-        for state in states:
-            if state.name not in self.states:
-                logger.debug("Registering state '{0}'".format(state.name))
-                self.add_state(state)
-
-        self.add_state("migrated")
-        self.add_ordered_transitions(loop=False)
+        logger = get_logger()
 
         if self._machine_pickle_file is None:
             logger.info("Execution started without machine pickle file")
@@ -54,12 +46,8 @@ def execute(
     def _execute(self, pickle_file: Optional[IO[Any]] = None) -> None:
         states_data = {}
 
-        # Read pickle file if it exists, to continue from the stored state
-        if pickle_file is not None and pickle_file.tell() > 0:
-            pickle_file.seek(0)
-
-            states_data = Unpickler(pickle_file).load()
-            self._restore_state_data(states_data)
+        if pickle_file is not None:
+            self._restore_state_data(pickle_file)
 
         try:
             while True:
@@ -110,7 +98,39 @@ def get_preflight_state_data(
     ) -> Any:
         return getattr(self.get_preflight_state(name), tag, default_value)
 
-    def _restore_state_data(self, data: Dict[str, Any]) -> None:
+    def register_states(self) -> None:
+        self._register_states(self._preflight_states + self._execution_states)
+
+        if self._machine_pickle_file is not None:
+            with open(self._machine_pickle_file, "rb") as file:
+                file.seek(1)
+                self._restore_state_data(file)
+
+    def _register_states(self, states: List[State]) -> None:
+        logger = get_logger()
+
+        for state in states:
+            if state.name not in self.states:
+                logger.debug("Registering state '{0}'".format(state.name))
+                self.add_state(state)
+
+        self.add_state("migrated")
+        self.add_ordered_transitions(loop=False)
+
+    """
+Read pickle file if it exists, to continue from the stored state. It is
+required that the position of the pointer of the pickle file given is not
+at the start.
+    """
+
+    def _restore_state_data(self, pickle_file: IO[Any]) -> None:
+        if pickle_file.tell() == 0:
+            return None
+
+        pickle_file.seek(0)
+
+        data = Unpickler(pickle_file).load()
+
         for state_name in data:
             try:
                 state = self.get_state(state_name)