Coverage for src/shephex/cli/slurm/profile.py: 41%
61 statements
« prev ^ index » next coverage.py v7.6.1, created at 2025-03-29 18:45 +0100
« prev ^ index » next coverage.py v7.6.1, created at 2025-03-29 18:45 +0100
1from pathlib import Path
3import rich
4import rich_click as click
5from rich.table import Table
7from shephex.cli.slurm.slurm import slurm
8from shephex.executor.slurm import SlurmExecutor, SlurmHeader, SlurmProfileManager
11@slurm.group()
12def profile() -> None:
13 """
14 Manage SLURM profiles.
15 """
16 pass
19@profile.command()
20@click.argument("path", required=False, type=click.Path(exists=True, file_okay=False, dir_okay=True))
21def directory(path: click.Path) -> None:
22 """
23 Set the default directory for Slurm profiles. If no path is provided, the current default directory is printed.
24 """
25 if path is None:
26 spm = SlurmProfileManager()
27 click.echo(spm.get_profile_directory())
28 return
30@profile.command()
31@click.argument("file", type=click.Path(exists=True, file_okay=True, dir_okay=False))
32@click.option("--name", "-n", help="Name of the profile. If not provided, the name of the file will be used.", default=None)
33@click.option("--overwrite", "-o", is_flag=True, help="Overwrite the profile if it already exists.", default=False)
34def add(file: click.Path, name: str, overwrite: bool) -> None:
35 """
36 Add a new SLURM profile.
37 """
38 import shutil
40 spm = SlurmProfileManager()
41 directory = spm.get_profile_directory()
43 if not file.endswith('.json'):
44 raise ValueError('Profile file must be a json file.')
46 if name is None:
47 name = Path(file).stem
49 new_path = (directory / name).with_suffix('.json')
51 if new_path.exists() and not overwrite:
52 raise FileExistsError(f'Profile {name} already exists in {directory}')
53 else:
54 shutil.copy(file, new_path)
56@profile.command(name="list")
57def list_profiles() -> None:
58 """
59 List all available SLURM profiles.
60 """
61 spm = SlurmProfileManager()
62 profiles = spm.get_all_profiles()
64 table = Table(title="SLURM Profiles")
65 table.add_column("Name")
66 table.add_column("Full Path")
68 for profile in profiles:
69 table.add_row(profile.stem, str(profile))
71 rich.print(table)
73@profile.command()
74@click.argument("name")
75@click.option("as_dict", "--dict", is_flag=True, help="Print the profile as a dictionary.", default=False)
76def print(name: str, as_dict: bool) -> None:
77 """
78 Print the contents of a SLURM profile.
79 """
80 spm = SlurmProfileManager()
81 profile = spm.get_profile(name)
82 if as_dict:
83 rich.print(profile)
84 return
86 executor = SlurmExecutor.from_profile(name)
87 rich.print(executor.header)
88 rich.print(executor._make_slurm_body([]))
92@profile.command()
93@click.argument("name")
94def delete(name: str) -> None:
95 """
96 Delete a SLURM profile.
97 """
98 spm = SlurmProfileManager()
99 profile = spm.get_profile_path(name)
100 profile.unlink()
101 click.echo(f'Profile {name} deleted.')