profiler.py

@@ -1,10 +1,63 @@
+#!/bin/python3
+
+from collections.abc import Iterable, Sequence
+from pathlib import Path
+from typing import Any
 import toml
 import subprocess
-import os
 import argparse
+from datetime import datetime
+import socket
+import os
+
+MACHINE_HOSTNAME = socket.gethostname()
+USERNAME = os.getlogin()
 
 
-def run_subprocess(command, verbose=False):
+class GenerateSystemAction(argparse.Action):
+    def __call__(
+        self,
+        _parser: argparse.ArgumentParser,
+        namespace: argparse.Namespace,
+        values: Any | None,
+        _option_string: str | None = None,
+    ) -> None:
+        if values is None:
+            # create auto generated name
+            now_time = datetime.now().strftime("%F_%k-%M-%S")
+            values = f"{MACHINE_HOSTNAME}_{USERNAME}_{now_time}.toml"
+
+        setattr(namespace, self.dest, values)
+
+
+parser = argparse.ArgumentParser(
+    description="Install packages and manage configurations from a TOML file."
+)
+parser.add_argument(
+    "--load", type=Path, help="Path to the TOML file to load and install packages."
+)
+parser.add_argument(
+    "--check", type=Path, help="Path to the TOML file to check for errors."
+)
+parser.add_argument(
+    "--generate-system",
+    type=str,
+    # allows to pass the argument without a value
+    nargs="?",
+    const=None,
+    action=GenerateSystemAction,
+    help="Generate a TOML file with system installed packages. pass the output file name or leave empty to auto generate.",
+)
+parser.add_argument(
+    "--generate-file",
+    type=str,
+    nargs=2,
+    metavar=("input", "output"),
+    help="Generate a TOML file from a text file. Provide input and output file paths.",
+)
+
+
+def run_subprocess(command: list[str], verbose=False):
     try:
         if verbose:
             result = subprocess.run(command)
@@ -19,21 +72,21 @@ def run_subprocess(command, verbose=False):
         print(f"Subprocess error: {e}")
 
 
-def install_generic_packages(command, packages):
+def install_generic_packages(command: list[str], packages: Iterable[str]):
     for package in packages:
         print(f"Installing package: {package}")
         run_subprocess(command + [package])
 
 
-def install_packages(packages):
+def install_packages(packages: Iterable[str]):
     install_generic_packages(["sudo", "pacman", "-S", "--noconfirm"], packages)
 
 
-def install_aur_packages(aur_packages):
+def install_aur_packages(aur_packages: Iterable[str]):
     install_generic_packages(["paru", "-S", "--noconfirm"], aur_packages)
 
 
-def install_flatpak_packages(flatpak_packages):
+def install_flatpak_packages(flatpak_packages: Iterable[str]):
     install_generic_packages(["flatpak", "install", "-y"], flatpak_packages)
 
 
@@ -46,14 +99,14 @@ def clone_config(config):
             run_subprocess(["git", "clone", url, target_dir])
 
 
-def enable_systemd_services(services):
+def enable_systemd_services(services: Iterable[str]):
     for service in services:
         print(f"Enabling systemd service: {service}")
         run_subprocess(["sudo", "systemctl", "enable", service])
         run_subprocess(["sudo", "systemctl", "start", service])
 
 
-def check_toml(toml_file):
+def check_toml(toml_file: Path):
     try:
         toml.load(toml_file)
         print(f"{toml_file} is valid.")
@@ -61,7 +114,7 @@ def check_toml(toml_file):
         print(f"Error in {toml_file}: {e}")
 
 
-def generate_system_toml(output_file):
+def generate_system_toml(output_file: Path | str):
     # Get Pacman packages
     pacman_packages = subprocess.check_output(
         ["pacman", "-Qeq"], text=True
@@ -102,7 +155,7 @@ def generate_system_toml(output_file):
     print(f"TOML file saved as {output_file}")
 
 
-def generate_file_toml(input_file, output_file):
+def generate_file_toml(input_file: Path | str, output_file: Path | str):
     config = {
         "packages": {"packages": []},
         "aur": {"aur_packages": []},
@@ -154,28 +207,6 @@ def generate_file_toml(input_file, output_file):
 
 
 def main():
-    parser = argparse.ArgumentParser(
-        description="Install packages and manage configurations from a TOML file."
-    )
-    parser.add_argument(
-        "--load", type=str, help="Path to the TOML file to load and install packages."
-    )
-    parser.add_argument(
-        "--check", type=str, help="Path to the TOML file to check for errors."
-    )
-    parser.add_argument(
-        "--generate-system",
-        type=str,
-        help="Generate a TOML file with system installed packages.",
-    )
-    parser.add_argument(
-        "--generate-file",
-        type=str,
-        nargs=2,
-        metavar=("input", "output"),
-        help="Generate a TOML file from a text file. Provide input and output file paths.",
-    )
-
     args = parser.parse_args()
 
     if args.load:
@@ -211,7 +242,12 @@ def main():
         check_toml(toml_file)
 
     elif args.generate_system:
-        output_file = args.generate_system
+        output_file: str = args.generate_system
+        output_file = (
+            output_file
+            if output_file.lower().endswith(".toml")
+            else output_file + ".toml"
+        )
         generate_system_toml(output_file)
 
     elif args.generate_file: