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

1from pathlib import Path 

2 

3import rich 

4import rich_click as click 

5from rich.table import Table 

6 

7from shephex.cli.slurm.slurm import slurm 

8from shephex.executor.slurm import SlurmExecutor, SlurmHeader, SlurmProfileManager 

9 

10 

11@slurm.group() 

12def profile() -> None: 

13 """ 

14 Manage SLURM profiles. 

15 """ 

16 pass 

17 

18 

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 

29 

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 

39 

40 spm = SlurmProfileManager() 

41 directory = spm.get_profile_directory() 

42 

43 if not file.endswith('.json'): 

44 raise ValueError('Profile file must be a json file.') 

45 

46 if name is None: 

47 name = Path(file).stem 

48 

49 new_path = (directory / name).with_suffix('.json') 

50 

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) 

55 

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() 

63 

64 table = Table(title="SLURM Profiles") 

65 table.add_column("Name") 

66 table.add_column("Full Path") 

67 

68 for profile in profiles: 

69 table.add_row(profile.stem, str(profile)) 

70 

71 rich.print(table) 

72 

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 

85 

86 executor = SlurmExecutor.from_profile(name) 

87 rich.print(executor.header) 

88 rich.print(executor._make_slurm_body([])) 

89 

90 

91 

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.') 

102 

103 

104 

105 

106 

107 

108 

109 

110 

111 

112 

113 

114 

115 

116 

117