diff --git a/src/ugrd/generator_helpers.py b/src/ugrd/generator_helpers.py index 2ef592c6..7b34323a 100644 --- a/src/ugrd/generator_helpers.py +++ b/src/ugrd/generator_helpers.py @@ -48,7 +48,7 @@ def _get_build_path(self, path: Path | str) -> Path: return get_subpath(build_dir, path) return get_subpath(get_subpath(self["tmpdir"], build_dir), path) - def _mkdir(self, path: Path, resolve_build=True) -> None: + def _mkdir(self, path: Path, resolve_build: bool = True) -> None: """ Creates a directory within the build directory. If resolve_build is True, the path is resolved to the build directory. @@ -78,7 +78,9 @@ def _mkdir(self, path: Path, resolve_build=True) -> None: else: self.logger.debug("Directory already exists: %s" % path_dir) - def _write(self, file_name: Path | str, contents: list[str] | str, chmod_mask=0o644, append=False) -> None: + def _write( + self, file_name: Path | str, contents: list[str] | str, chmod_mask: int = 0o644, append: bool = False + ) -> None: """ Writes test to a file within the build directory. Sets the passed chmod_mask. @@ -121,7 +123,7 @@ def _write(self, file_name: Path | str, contents: list[str] | str, chmod_mask=0o file_path.chmod(chmod_mask) self.logger.debug("[%s] Set file permissions: %s" % (file_path, chmod_mask)) - def _copy(self, source: Path | str, dest=None) -> None: + def _copy(self, source: Path | str, dest: Path | str | None = None) -> None: """Copies a file into the initramfs build directory. If a destination is not provided, the source is used, under the build directory. @@ -209,14 +211,16 @@ def _symlink(self, source: Path | str, target: Path | str) -> None: ) target.symlink_to(source) - def _run(self, args: list[str], timeout=None, fail_silent=False, fail_hard=True) -> CompletedProcess: + def _run( + self, args: list[str], timeout: int | None = None, fail_silent: bool = False, fail_hard: bool = True + ) -> CompletedProcess: # type: ignore[type-arg] """Runs a command, returns the CompletedProcess object on success. If a timeout is set, the command will fail hard if it times out. If fail_silent is set, non-zero return codes will not log stderr/stdout. If fail_hard is set, non-zero return codes will raise a RuntimeError. """ - def print_err(ret) -> None: + def print_err(ret: CompletedProcess | TimeoutExpired) -> None: # type: ignore[type-arg] if args := ret.args: if isinstance(args, tuple): args = args[0] # When there's a timeout, args is a (args, timeout) tuple @@ -244,7 +248,7 @@ def print_err(ret) -> None: return cmd - def _rotate_old(self, file_name: Path, sequence=0) -> None: + def _rotate_old(self, file_name: Path, sequence: int = 0) -> None: """Copies a file to file_name.old then file_name.old.n, where n is the next number in the sequence""" # Nothing to do if the file doesn't exist if not file_name.is_file(): @@ -308,7 +312,7 @@ def sort_hook_functions(self, hook: str) -> None: if not before and not after: return self.logger.debug("No import order specified for hook: %s" % hook) - def iter_order(order, direction): + def iter_order(order: dict[str, list[str]], direction: str) -> bool: """Iterate over all functions in an import order list, using this information to move the order of function names in the import list. @@ -325,7 +329,7 @@ def iter_order(order, direction): continue assert other_index >= 0, "Function not found in import list: %s" % other_func - def reorder_func(direction): + def reorder_func(direction: str) -> None: """Reorders the function based on the direction.""" if direction == "before": # Move the function before the other function self.logger.debug( diff --git a/src/ugrd/initramfs_config.py b/src/ugrd/initramfs_config.py index 19ce05b6..b3c9ce8d 100644 --- a/src/ugrd/initramfs_config.py +++ b/src/ugrd/initramfs_config.py @@ -7,6 +7,7 @@ from pathlib import Path from queue import Queue from typing import Any +from types import ModuleType from pycpio import PyCPIO from zenlib.logging import LoggerMixIn @@ -48,11 +49,11 @@ class InitramfsConfig(LoggerMixIn, UserDict): def __init__( self, - startup_args: dict | None = None, + startup_args: dict[str, Any] | None = None, config_file: Path | str | None = None, NO_BASE: bool = False, - *args, - **kwargs, + *args: Any, + **kwargs: Any, ) -> None: """Initialize the initramfs config @@ -120,7 +121,7 @@ def _enqueue(self, key: str, value: Any) -> None: f"[{c_(key, 'blue', background=True)}] Adding parameter to processing queue: {c_(value, 'yellow')}" ) - def __setitem__(self, key: str, value) -> None: + def __setitem__(self, key: str, value: Any) -> None: """Custom setitem for the config dict If the dict is validated, refuse to set config If it's the final stage and it's not validated, raise a critical warning @@ -132,13 +133,13 @@ def __setitem__(self, key: str, value) -> None: For everything but the logger, queue values if they are not registered """ if self["validated"]: - return self.logger.error( - f"[{c_(key, 'yellow')}] Config is validated, refusing to set value: {c_(value, 'red')}" - ) + self.logger.error(f"[{c_(key, 'yellow')}] Config is validated, refusing to set value: {c_(value, 'red')}") + return if self["stage"] == "final" and not self["validated"]: - return self.logger.critical( + self.logger.critical( f"[{c_(key, 'yellow')}] Config is finalized but invalid, refusing to set value: {c_(value, 'red')}" ) + return if not self._check_late(key): return self._enqueue(key, value) @@ -156,7 +157,7 @@ def __setitem__(self, key: str, value) -> None: if key != "logger": self._enqueue(key, value) - def handle_parameter(self, key: str, value) -> None: + def handle_parameter(self, key: str, value: Any) -> None: """ Handles a config parameter, setting the value and processing it if the type is known. Raises a KeyError if the parameter is not registered. @@ -171,14 +172,16 @@ def handle_parameter(self, key: str, value) -> None: if expected_type: if expected_type.__name__ == "InitramfsGenerator": self.data[key] = value - return self.logger.debug(f"Setting InitramfsGenerator: {c_(key, 'magenta', bold=True)}") + self.logger.debug(f"Setting InitramfsGenerator: {c_(key, 'magenta', bold=True)}") + return break # Break and raise an exception if the type is not found else: raise KeyError(f"Parameter not registered: {c_(key, 'red')}") if hasattr(self, f"_process_{key}"): # The builtin function is decorated and can handle plural self.logger.log(5, f"[{c_(key, 'blue')}] Using builtin setitem: _process_{key}") - return getattr(self, f"_process_{key}")(value) + getattr(self, f"_process_{key}")(value) + return # Don't use masked processing functions for custom values, fall back to standard setters def check_mask(import_name: str) -> bool: @@ -192,7 +195,8 @@ def check_mask(import_name: str) -> bool: self.logger.log( 5, f"[{c_(key, 'blue')}] Using custom setitem: {c_(func.__name__, 'blue', underline=True)}" ) - return func(self, value) + func(self, value) + return if func := self["custom_processing"].get(f"_process_{key}_multi"): if check_mask(func.__name__): @@ -202,18 +206,22 @@ def check_mask(import_name: str) -> bool: 5, f"[{c_(key, 'blue')}] Using custom plural setitem: {c_(func.__name__, 'blue', underline=True, bold=True)}", ) - return handle_plural(func)(self, value) + handle_plural(func)(self, value) + return if expected_type in (list, NoDupFlatList): # Append to lists, don't replace self.logger.log(5, f"[{c_(key, 'blue')}] Using list setitem") - return self[key].append(value) + self[key].append(value) + return if expected_type is dict: # Create new keys, update existing if key not in self: self.logger.log(5, f"[{c_(key, 'blue')}] Setting dict to: {value}") - return super().__setitem__(key, value) + super().__setitem__(key, value) + return self.logger.log(5, f"[{c_(key, 'blue')}] Updating dict with: {value}") - return self[key].update(value) + self[key].update(value) + return casted_value = expected_type(value) self.logger.debug( @@ -285,35 +293,35 @@ def _process_unprocessed(self, parameter_name: str) -> None: ) self[parameter_name] = value - def _process_import_order(self, import_order: dict) -> None: + def _process_import_order(self, import_order: dict[str, dict[str, list[str] | str]]) -> None: """Processes the import order, setting the order requirements for import functions. Ensures the order type is valid (before, after), that the function is not ordered after itself. Ensures that the same function/target is not in another order type. """ - self.logger.debug("Processing import order:\n%s" % pretty_print(import_order)) + self.logger.debug(f"Processing import order:\n{pretty_print(import_order)}") order_types = ["before", "after"] for order_type, order_dict in import_order.items(): if order_type not in order_types: - raise ValueError("Invalid import order type: %s" % order_type) + raise ValueError(f"Invalid import order type: {c_(order_type, 'red')}") for function in order_dict: targets = order_dict[function] if not isinstance(targets, list): targets = [targets] if function in targets: - raise ValueError("Function cannot be ordered after itself: %s" % function) + raise ValueError(f"Function cannot be ordered after itself: {c_(function, 'red')}") for other_target in [self["import_order"].get(ot, {}) for ot in order_types if ot != order_type]: if function in other_target and any(target in other_target[function] for target in targets): - raise ValueError("Function cannot be ordered in multiple types: %s" % function) + raise ValueError(f"Function cannot be ordered in multiple types: {c_(function, 'red')}") order_dict[function] = targets if order_type not in self["import_order"]: self["import_order"][order_type] = {} self["import_order"][order_type].update(order_dict) - self.logger.debug("Registered import order requirements: %s" % import_order) + self.logger.debug(f"Registered import order requirements:\n{import_order}") - def _import_external_module(self, module_name: str): + def _import_external_module(self, module_name: str) -> ModuleType: """Given a module name, attempts to load it from /var/lib/ugrd, returning the module""" module_path = Path("/var/lib/ugrd/" + module_name.replace(".", "/")).with_suffix(".py") self.logger.debug(f"Attempting to sideload module from: {c_(module_path, 'green')}") @@ -349,7 +357,7 @@ def _import_external_module(self, module_name: str): return module @handle_plural - def _process_imports(self, import_type: str, import_value: dict) -> None: + def _process_imports(self, import_type: str, import_value: dict[str, list[str]]) -> None: """Processes imports in a module, importing the functions and adding them to the appropriate list.""" for module_name, function_names in import_value.items(): f_name = f"[{c_(module_name, 'green')}]({c_(import_type, underline=True)})" @@ -538,7 +546,8 @@ def _validate(self) -> None: raise ValidationError( f"Failed to validate config. Unprocessed values: {c_(unprocessed_values, 'red', bold=True)}" ) - return self.logger.critical(f"Unprocessed config values: {c_(unprocessed_values, 'red', bold=True)}") + self.logger.critical(f"Unprocessed config values: {c_(unprocessed_values, 'red', bold=True)}") + return self.data["validated"] = True def __str__(self) -> str: diff --git a/src/ugrd/initramfs_generator.py b/src/ugrd/initramfs_generator.py index 5a52b2da..1ec72270 100644 --- a/src/ugrd/initramfs_generator.py +++ b/src/ugrd/initramfs_generator.py @@ -1,7 +1,7 @@ from importlib.metadata import version from pathlib import Path from textwrap import dedent -from typing import Any +from typing import Any, Callable from zenlib.logging import LoggerMixIn from zenlib.util import colorize as c_ @@ -15,7 +15,7 @@ class InitramfsGenerator(GeneratorHelpers, LoggerMixIn): - def __init__(self, config: Path | str | None = DEFAULT_CONFIG_PATH, *args, **kwargs) -> None: + def __init__(self, config: Path | str | None = DEFAULT_CONFIG_PATH, *args: Any, **kwargs: Any) -> None: self.init_logger(args, kwargs) self.config_dict = InitramfsConfig( NO_BASE=kwargs.pop("NO_BASE", False), logger=self.logger, startup_args=kwargs, config_file=config @@ -32,19 +32,19 @@ def __init__(self, config: Path | str | None = DEFAULT_CONFIG_PATH, *args, **kwa self.init_types = ["init_debug", "init_main", "init_mount"] # If the initramfs generator is used as a dictionary, it will use the config_dict. - def __setitem__(self, key, value) -> None: + def __setitem__(self, key: str, value: Any) -> None: self.config_dict[key] = value - def __getitem__(self, item) -> Any: + def __getitem__(self, item: str) -> Any: return self.config_dict[item] - def __contains__(self, item) -> bool: + def __contains__(self, item: str) -> bool: return item in self.config_dict - def get(self, item, default=None) -> Any: + def get(self, item: str, default: Any = None) -> Any: return self.config_dict.get(item, default) - def __getattr__(self, item) -> Any: + def __getattr__(self, item: str) -> Any: """Allows access to the config dict via the InitramfsGenerator object.""" if item not in self.__dict__ and item != "config_dict": return self[item] @@ -62,7 +62,9 @@ def build(self) -> None: self.run_checks() self.run_tests() - def run_func(self, function, force_include=False, force_exclude=False) -> list[str] | None: + def run_func( + self, function: Callable[..., list[str] | str | None], force_include: bool = False, force_exclude: bool = False + ) -> list[str] | None: """ Runs an imported function. The function should return str | list[str] | none @@ -109,12 +111,12 @@ def run_func(self, function, force_include=False, force_exclude=False) -> list[s self.logger.debug(f"Created function alias: {c_(function.__name__, 'blue')}") return [function.__name__] elif force_include: - raise ValueError(f"Force included function returned no output:{c_(function.__name_, 'red')}") + raise ValueError(f"Force included function returned no output:{c_(function.__name__, 'red')}") else: self.logger.debug(f"Function returned no output: {c_(function.__name__, 'yellow')}") return None - def run_hook(self, hook: str, *args, **kwargs) -> list[str]: + def run_hook(self, hook: str, *args: Any, **kwargs: Any) -> list[str]: """Runs all functions for the specified hook. If the function is masked, it will be skipped. If the function is in import_order, handle the ordering @@ -267,7 +269,7 @@ def run_init_hook(self, level: str) -> list[str]: self.logger.debug("No output for init level: %s" % level) return [] - def _log_run(self, logline) -> None: + def _log_run(self, logline: str) -> None: self.logger.info(f"-- | {c_(logline, 'blue', bold=True)}") def __str__(self) -> str: diff --git a/src/ugrd/initramfs_protocol.py b/src/ugrd/initramfs_protocol.py index c3a8f535..0e2efcd4 100644 --- a/src/ugrd/initramfs_protocol.py +++ b/src/ugrd/initramfs_protocol.py @@ -1,6 +1,6 @@ -from typing import Any, Protocol from pathlib import Path from subprocess import CompletedProcess +from typing import Any, Protocol from zenlib.typing import HasLogger @@ -20,11 +20,14 @@ def __setitem__(self, item: str, value: Any) -> None: ... def __contains__(self, key: str) -> bool: ... # Add definitions for helper functions - def _mkdir(self, path: Path, resolve_build=True) -> None: ... - def _write(self, file_name: Path | str, contents: list[str] | str, chmod_mask=0o644, append=False) -> None: ... - def _copy(self, source: Path | str, dest=None) -> None: ... + def _mkdir(self, path: Path, resolve_build: bool = True) -> None: ... + def _write( + self, file_name: Path | str, contents: list[str] | str, chmod_mask: int = 0o644, append: bool = False + ) -> None: ... + def _copy(self, source: Path | str, dest: Path | str | None = None) -> None: ... def _symlink(self, source: Path | str, target: Path | str) -> None: ... - def _run(self, args: list[str], timeout=None, fail_silent=False, fail_hard=True) -> CompletedProcess: ... - def _rotate_old(self, file_name: Path, sequence=0) -> None: ... + def _run( + self, args: list[str], timeout: int | None = None, fail_silent: bool = False, fail_hard: bool = True + ) -> CompletedProcess: ... # type: ignore[type-arg] + def _rotate_old(self, file_name: Path, sequence: int = 0) -> None: ... def sort_hook_functions(self, hook: str) -> None: ... -