Coverage for src/shephex/cli/slurm/profile.py: 100%

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, SlurmProfileManager 

9 

10 

11@slurm.group() 

12def profile() -> None: 

13 """ 

14 Manage SLURM profiles. 

15 """ 

16 ... # pragma: no cover 

17 

18 

19@profile.command() 

20@click.argument( 

21 'path', required=False, type=click.Path(exists=True, file_okay=False, dir_okay=True) 

22) 

23def directory(path: click.Path) -> None: 

24 """ 

25 Set the default directory for Slurm profiles. If no path is provided, the current default directory is printed. 

26 """ 

27 if path is None: 

28 spm = SlurmProfileManager() 

29 click.echo(spm.get_profile_directory()) 

30 return 

31 

32 

33@profile.command() 

34@click.argument( 

35 'file_name', 

36 type=click.Path(exists=True, file_okay=True, dir_okay=False, path_type=Path), 

37) 

38@click.option( 

39 '--name', 

40 '-n', 

41 help='Name of the profile. If not provided, the name of the file will be used.', 

42 default=None, 

43) 

44@click.option( 

45 '--overwrite', 

46 '-o', 

47 is_flag=True, 

48 help='Overwrite the profile if it already exists.', 

49 default=False, 

50) 

51def add(file_name: Path, name: str, overwrite: bool) -> None: 

52 """ 

53 Add a new SLURM profile. 

54 """ 

55 spm = SlurmProfileManager() 

56 spm.add_profile(file_name, name, overwrite) 

57 

58 

59@profile.command(name='list') 

60def list_profiles() -> None: 

61 """ 

62 List all available SLURM profiles. 

63 """ 

64 spm = SlurmProfileManager() 

65 profiles = spm.get_all_profiles() 

66 

67 table = Table(title='SLURM Profiles') 

68 table.add_column('Name') 

69 table.add_column('Full Path') 

70 

71 for profile in profiles: 

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

73 

74 rich.print(table) 

75 

76 

77@profile.command(name='print') 

78@click.argument('name') 

79@click.option( 

80 'as_dict', 

81 '--dict', 

82 is_flag=True, 

83 help='Print the profile as a dictionary.', 

84 default=False, 

85) 

86def print_profile(name: str, as_dict: bool) -> None: 

87 """ 

88 Print the contents of a SLURM profile. 

89 """ 

90 spm = SlurmProfileManager() 

91 profile = spm.get_profile(name) 

92 if as_dict: 

93 rich.print(profile) 

94 return 

95 

96 executor = SlurmExecutor.from_profile(name, safety_check=False) 

97 rich.print(executor.header) 

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

99 

100 

101@profile.command() 

102@click.argument('name') 

103def delete(name: str) -> None: 

104 """ 

105 Delete a SLURM profile. 

106 """ 

107 spm = SlurmProfileManager() 

108 profile = spm.get_profile_path(name) 

109 profile.unlink() 

110 click.echo(f'Profile {name} deleted.')