Merge branch 'extend-status-cmd' into 'main'

add status command arguments

Closes tickets#27

See merge request !118

spack-environments @ 91ba21f4

+Subproject commit 91ba21f40ff9e635da15a4c154360f5e969d6fe4

src/mpsd_software_manager/mpsd_software.py

@@ -16,6 +16,7 @@ import time
 from functools import cache
 from pathlib import Path
 from typing import List, Tuple, Union
+from rich import print as rprint
 
 __version__ = importlib.metadata.version(__package__ or __name__)
 
@@ -1259,8 +1260,36 @@ def start_new_environment(release, from_release, target_dir):
     raise NotImplementedError(msg)
 
 
-def environment_status(mpsd_release: str, root_dir: Path) -> Union[dict, None]:
-    """Show status of release in installation.
+def list_installed_releases(root_dir: Path, print_output: bool = False) -> List[str]:
+    """
+    List installed releases.
+
+    Parameters
+    ----------
+    root_dir : pathlib.Path
+
+    Returns
+    -------
+    installed_releases : list
+        A list of strings representing the installed releases.
+    """
+    plog = logging.getLogger("print")
+    list_of_files = os.listdir(root_dir)
+    installed_releases = [
+        x for x in list_of_files if (root_dir / x / "spack-environments").exists()
+    ]
+    if print_output:
+        plog.info("Available MPSD software releases:")
+        for release in installed_releases:
+            plog.info(f"    {release}")
+    return installed_releases
+
+
+def list_installed_toolchains(
+    mpsd_release: str, root_dir: Path, print_output: bool = False
+) -> Union[dict, None]:
+    """
+    List installed toolchains.
 
     Parameters
     ----------
@@ -1270,6 +1299,8 @@ def environment_status(mpsd_release: str, root_dir: Path) -> Union[dict, None]:
         A Path object pointing to the root directory of the installation.
         Expect a subfolder root/mpsd_release in which we search for the
         toolchains.
+    print_output : bool, optional
+        A boolean indicating whether to print the output to the terminal.
 
     Returns
     -------
@@ -1280,10 +1311,7 @@ def environment_status(mpsd_release: str, root_dir: Path) -> Union[dict, None]:
 
         Note: only toolchains can be reported at the moment (i.e. package_sets
         such as global and global_generic are missing, even if installed).
-
     """
-    msg = f"Showing status of release {mpsd_release} in {root_dir}"
-    logging.info(msg)
     plog = logging.getLogger("print")
     release_base_dir = root_dir / mpsd_release
     microarch = get_native_microarchitecture()
@@ -1334,16 +1362,182 @@ def environment_status(mpsd_release: str, root_dir: Path) -> Union[dict, None]:
 
     # pretty print the toolchain map key as the heading
     # and the value as the list of toolchains
-    plog.info(f"Installed toolchains ({mpsd_release}):\n")
-    for microarch, toolchains in toolchain_map.items():
-        plog.info(f"- {microarch}")
-        for toolchain in toolchains:
-            plog.info(f"    {toolchain}")
-        plog.info(f"    [module use {str(release_base_dir / microarch / 'lmod/Core')}]")
-        plog.info("")
+    if print_output:
+        plog.info(f"Installed toolchains ({mpsd_release}):\n")
+        for microarch, toolchains in toolchain_map.items():
+            plog.info(f"- {microarch}")
+            for toolchain in toolchains:
+                plog.info(f"    {toolchain}")
+            plog.info(
+                f"    [module use {str(release_base_dir / microarch / 'lmod/Core')}]"
+            )
+            plog.info("")
     return toolchain_map
 
 
+def pretty_print_spec(spec: str) -> None:
+    """
+    Print the specs with colours using rich.
+
+    - packages in white (everything until first %)
+    - compiler in green (everything between % and first+)
+    - variants in cyan (everything that starts with +)
+    - build_system in yellow (everything that starts with build_system=)
+    - architecture in purple (everything that starts with arch=)
+    """
+    # Note that this implementation necessitates the definition of
+    # flags in the order in which we ask spack to format the output
+    # also for flags that need the same colour because they are
+    # interchangeable (like `+` and `~`) we need to define them together
+    colour_map = {
+        "%": "green",
+        "+": "cyan",
+        "~": "cyan",
+        "build_system=": "yellow",
+        "libs=": "blue",
+        "arch=": "purple",
+    }
+
+    prev_colour = ""
+
+    for flag in colour_map.keys():
+        # If the flag is in the spec string,
+        # replace it with: previous closing colour, new colour, flag
+        if flag in spec:
+            if (
+                colour_map[flag] not in prev_colour
+            ):  # avoid duplicates for eg when having both ~ and +
+                spec = spec.replace(flag, f"{prev_colour}[{colour_map[flag]}]{flag}", 1)
+                prev_colour = f"[/{colour_map[flag]}]"  # for next iter
+
+    # Add the final closing tag to the spec string
+    spec += prev_colour
+    rprint(spec)
+
+
+def list_installed_packages(
+    mpsd_release: str, root_dir: Path, package_set: str, microarch: str
+) -> Union[List[str], None]:
+    """
+    List installed packages and their specs.
+
+    Uses `spack -e package_set find` to list the installed packages,
+    in the following format
+    "{name}{@versions}{%compiler.name}{@compiler.versions}{compiler_flags}{variants}{arch=architecture}"
+
+    Parameters
+    ----------
+    mpsd_release : str
+        A string representing the MPSD release version.
+    root_dir : pathlib.Path
+        A Path object pointing to the root directory of the installation.
+        Expect a subfolder root/mpsd_release in which we search for the
+        toolchains.
+    package_set : str
+        A string representing the package_sets to show the packages for.
+    microarch : str
+        A string representing the microarchitecture to show the packages for.
+
+    Returns
+    -------
+    list
+        A list of strings representing the packages installed for the
+        specified package_sets and microarch.
+        If the release is not installed/found, None is returned.
+
+    """
+    plog = logging.getLogger("print")
+    plog.info(f"listing packages installed for {package_set=}, {microarch=}")
+    spack_dir = root_dir / mpsd_release / microarch / "spack"
+    spack_env = spack_dir / "share" / "spack" / "setup-env.sh"
+    commands_to_execute = [
+        f"export SPACK_ROOT={spack_dir}",  # need to set SPACK_ROOT in dash and sh
+        f". {spack_env}",
+        f"spack -e {package_set}"
+        " find --format "
+        r"{name}{@versions}{%compiler.name}{@compiler.versions}{compiler_flags}{variants}{arch=architecture}",
+    ]
+    process = run(
+        " && ".join(commands_to_execute), shell=True, check=True, capture_output=True
+    )
+    package_list = process.stdout.decode().strip().split("\n")
+    for package in package_list:
+        pretty_print_spec(package)
+    return package_list
+
+
+def environment_status(
+    mpsd_release: str, root_dir: Path, package_set="NONE"
+) -> Union[dict, List[str], None]:
+    """Show status of release in installation.
+
+    - 1) If no mpsd_release, list available releases
+    - 2) If mpsd_release, list available toolchains
+    - 3) If mpsd_release and toolchain, list available packages
+
+    Parameters
+    ----------
+    mpsd_release : str
+        A string representing the MPSD release version.
+    root_dir : pathlib.Path
+        A Path object pointing to the root directory of the installation.
+        Expect a subfolder root/mpsd_release in which we search for the
+        toolchains.
+    package_set : str, optional
+        A string representing the package_sets to show the status for.
+
+    Returns
+    -------
+    installed_release : List[str]
+        A list of installed (valid) releases.
+
+    OR
+    toolchain_map : dict
+        A dictionary containing available microarchitectures as keys and
+        a list of available package_sets as values for each microarchitecture.
+        If the release is not installed/found, None is returned.
+
+        Note: only toolchains can be reported at the moment (i.e. package_sets
+        such as global and global_generic are missing, even if installed).
+
+    OR
+    package_list : List[str]
+        A list of strings representing the packages installed for the
+        specified package_sets and microarch.
+
+    """
+    msg = f"Showing status of release {mpsd_release} in {root_dir}"
+    logging.info(msg)
+    if not mpsd_release:
+        # 1) if no mpsd_release is specified, list available releases
+        return list_installed_releases(root_dir=root_dir, print_output=True)
+    # 2) if mpsd_release is specified, list installed toolchains
+    # Test is the mpsd_release is valid
+    if mpsd_release not in list_installed_releases(root_dir=root_dir):
+        logging.error(f"MPSD release '{mpsd_release}' is not available.")
+        return None
+    if package_set == "NONE":
+        return list_installed_toolchains(
+            mpsd_release=mpsd_release, root_dir=root_dir, print_output=True
+        )
+    # 3) if mpsd_release and toolchain is specified, list installed packages
+    # check that the package-set is a valid toolchain
+    if (
+        package_set
+        not in list_installed_toolchains(mpsd_release=mpsd_release, root_dir=root_dir)[
+            get_native_microarchitecture()
+        ]
+    ):
+        logging.error(f"Package-set '{package_set}' is not available.")
+        return None
+    return list_installed_packages(
+        mpsd_release=mpsd_release,
+        root_dir=root_dir,
+        package_set=package_set,
+        microarch=get_native_microarchitecture(),
+    )
+
+
 def initialise_environment(root_dir: Path) -> None:
     """Initialize the software environment.
 
@@ -1478,13 +1672,13 @@ def main():
 
         else:
             # most commands except need a release version
-            if cmd in ["install", "prepare", "reinstall", "remove", "status"]:
+            if cmd in ["install", "prepare", "reinstall", "remove"]:
                 subp.add_argument(
                     "release",
                     type=str,
                     help="Release version to prepare, install, reinstall or remove",
                 )
-            elif cmd in ["available"]:
+            elif cmd in ["available", "status"]:
                 # for some commands the release version is optional
                 subp.add_argument(
                     "release",
@@ -1508,6 +1702,7 @@ def main():
                     default="NONE",
                     help=package_set_help,
                 )
+                # TODO Move the enable-build-cache flag to only 'install' cmd
                 subp.add_argument(
                     "--enable-build-cache",
                     action="store_true",
@@ -1516,6 +1711,14 @@ def main():
                         "consumes time and disk space."
                     ),
                 )
+            if cmd in ["status"]:
+                subp.add_argument(
+                    "package_set",
+                    type=str,
+                    nargs="?",
+                    default="NONE",
+                    help="Package set to show status for.",
+                )
     # Carry out the action
     args = parser.parse_args()
 
@@ -1558,8 +1761,10 @@ def main():
         args.loglevel,
         apex_log_file,
     )
-
-    record_script_execution_summary(root_dir, apex_log_file)
+    if args.action not in ["status", "available"]:
+        # record the script execution summary only if
+        # the action is one that changes files on disk
+        record_script_execution_summary(root_dir, apex_log_file)
     # Check the command and run related function
     if args.action == "remove":
         remove_environment(args.release, root_dir, args.package_set)
@@ -1570,7 +1775,7 @@ def main():
             args.release, args.package_set, root_dir, args.enable_build_cache
         )
     elif args.action == "status":
-        _ = environment_status(args.release, root_dir)
+        environment_status(args.release, root_dir, args.package_set)
     elif args.action == "prepare":
         prepare_environment(args.release, root_dir)
     elif args.action == "available":

tests/test_mpsd_software.py

@@ -482,6 +482,8 @@ def create_fake_environment(tmp_path, mpsd_release, expected_toolchain_map=None)
         test_microarch = mod.get_native_microarchitecture()
         expected_toolchain_map = {test_microarch: ["foss2021a", "intel2021a"]}
 
+        spe_folder = tmp_path / mpsd_release / "spack-environments"
+        spe_folder.mkdir(parents=True, exist_ok=True)
     for microarch in expected_toolchain_map.keys():
         toolchain_lmod_folder = (
             tmp_path / mpsd_release / microarch / "lmod" / "Core" / "toolchains"
@@ -501,18 +503,68 @@ def create_fake_environment(tmp_path, mpsd_release, expected_toolchain_map=None)
     return expected_toolchain_map
 
 
-def test_environment_status(tmp_path):
-    """Test that the environment status is correct."""
+def check_for_valid_spec_syntax(spec: str, package_name: str):
+    """Check if the spec is valid.
+
+    Assuming the format of the spec as:
+    {name}{@versions}{%compiler.name}{@compiler.versions}{compiler_flags}{variants}{arch=architecture}
+    we ensure that:
+    - package name is correct
+    - there are atleast 2 @ symbols (version for package and compiler)
+    - there is atleast 1 % symbol  (compiler specification)
+    - there is atleast 1 arch=
+    """
+    assert spec.count("@") >= 2
+    assert spec.count("%") >= 1
+    assert "build_system=" in spec
+    assert "arch=" in spec
+    assert spec.split("@")[0] == package_name
+
+
+def test_environment_status(tmp_path, simple_toolchain):
+    """Test that the environment status is correct.
+
+    The status command has the following three usage:
+
+    - 1) If no mpsd_release, list available releases
+    - 2) If mpsd_release, list available toolchains
+    - 3) If mpsd_release and toolchain, list available packages
+
+    We need to test all the three cases.
+    """
+
+    # 1) If no mpsd_release, list available releases
+    list_of_release_in_empty_dir = mod.environment_status(None, tmp_path)
+    assert list_of_release_in_empty_dir == []
+    (tmp_path / "test_case1" / "dev-23a" / "spack-environments").mkdir(
+        parents=True, exist_ok=True
+    )
+    (tmp_path / "test_case1" / "fake_release").mkdir(parents=True, exist_ok=True)
+    list_of_release = mod.environment_status(None, tmp_path / "test_case1")
+    assert list_of_release == ["dev-23a"]
+
+    # 2) If mpsd_release, list available toolchains
     toolchain_map = mod.environment_status("fake-release", tmp_path)
     assert toolchain_map is None
     mpsd_release = "dev-23a"
     expected_toolchain_map = create_fake_environment(tmp_path, mpsd_release)
-    # check that the environment statuxis is correct
+    # check that the environment status is is correct
     toolchain_map = mod.environment_status(mpsd_release, tmp_path)
     # convert each list to a set to ensure that the order doesn't matter
     for microarch in expected_toolchain_map.keys():
         assert set(toolchain_map[microarch]) == set(expected_toolchain_map[microarch])
 
+    # 3) If mpsd_release and toolchain, list available packages
+    install_test_release(tmp_path / "test_case3", simple_toolchain)
+    package_list = mod.environment_status(
+        mpsd_release, tmp_path / "test_case3", "toolchain2"
+    )
+    assert len(package_list) == 2  # we installed zlib and zstd only
+    check_for_valid_spec_syntax(
+        package_list[0], "zlib"
+    )  # the list is always in alphabetical order
+    check_for_valid_spec_syntax(package_list[1], "zstd")
+
 
 def test_initialise_environment(tmp_path):
     """Test that init_file is created as expected."""