Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 12 additions & 8 deletions src/ugrd/generator_helpers.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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.

Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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():
Expand Down Expand Up @@ -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.

Expand All @@ -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(
Expand Down
59 changes: 34 additions & 25 deletions src/ugrd/initramfs_config.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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

Expand Down Expand Up @@ -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
Expand All @@ -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)
Expand All @@ -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.
Expand All @@ -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:
Expand All @@ -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__):
Expand All @@ -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(
Expand Down Expand Up @@ -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')}")
Expand Down Expand Up @@ -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)})"
Expand Down Expand Up @@ -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:
Expand Down
24 changes: 13 additions & 11 deletions src/ugrd/initramfs_generator.py
Original file line number Diff line number Diff line change
@@ -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_
Expand All @@ -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
Expand All @@ -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]
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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:
Expand Down
17 changes: 10 additions & 7 deletions src/ugrd/initramfs_protocol.py
Original file line number Diff line number Diff line change
@@ -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

Expand All @@ -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: ...

Loading