Coverage for src/shephex/executor/slurm/slurm_body.py: 100%

20 statements  

« prev     ^ index     » next       coverage.py v7.6.1, created at 2025-03-29 18:45 +0100

1from typing import List, Self, Union 

2 

3 

4class SlurmBody: 

5 def __init__(self, commands: Union[List, str] = None) -> None: 

6 """ 

7 Body portion of a Slurm script. That is the commands that will be executed, 

8 e.g. 

9 "python my_script.py" 

10 

11 Note these are generally not checked for correctness. 

12 

13 Parameters 

14 ---------- 

15 commands : Union[List, str], optional 

16 List of commands to be executed, by default None 

17 """ 

18 self.commands = [] 

19 

20 if commands is not None: 

21 self.add(commands) 

22 

23 def add(self, command: Union[List[str], str]) -> None: 

24 """ 

25 Add a command to the body 

26 

27 Parameters 

28 ---------- 

29 command : Union[List, str] 

30 Command to be added 

31 """ 

32 if isinstance(command, list): 

33 if not all(isinstance(c, str) for c in command): 

34 raise TypeError('All commands must be strings') 

35 self.commands.extend(command) 

36 else: 

37 if not isinstance(command, str): 

38 raise TypeError('Command must be a string') 

39 self.commands.append(command) 

40 

41 def __repr__(self) -> str: 

42 return '\n'.join(self.commands) 

43 

44 def __add__(self, other: Self) -> Self: 

45 """ 

46 Concatenate two SlurmBody objects. Order of operations matters! 

47 

48 Parameters 

49 ---------- 

50 other : Self 

51 SlurmBody object to concatenate with. 

52 """ 

53 body = SlurmBody() 

54 body.commands = self.commands + other.commands 

55 return body