Coverage for ezdag/options.py: 99.0%

96 statements  

« prev     ^ index     » next       coverage.py v7.6.4, created at 2024-10-29 15:59 -0700

1# Copyright (C) 2020 Patrick Godwin 

2# 

3# This Source Code Form is subject to the terms of the Mozilla Public License, v2.0. 

4# If a copy of the MPL was not distributed with this file, You can obtain one at 

5# <https://mozilla.org/MPL/2.0/>. 

6# 

7# SPDX-License-Identifier: MPL-2.0 

8 

9from __future__ import annotations 

10 

11from collections.abc import Iterable 

12from dataclasses import dataclass, field 

13from typing import Callable, ClassVar 

14 

15from . import path 

16 

17PROTECTED_CONDOR_VARS = {"input", "output", "rootdir"} 

18 

19 

20@dataclass 

21class Argument: 

22 """Defines a command-line argument (positional). 

23 

24 This provides some extra functionality over defining command line 

25 argument explicitly, in addition to some extra parameters which 

26 sets how condor interprets how to handle them within the DAG 

27 and within submit descriptions. 

28 

29 Parameters 

30 ---------- 

31 name 

32 The option name. Since this is a positional argument, it is not 

33 used explicitly in the command, but is needed to define 

34 variable names within jobs. 

35 argument 

36 The positional argument value(s) used in a command. 

37 track 

38 Whether to track files defined here and used externally within 

39 jobs to determine parent-child relationships when nodes specify 

40 this option as an input or output. On by default. 

41 suppress 

42 Whether to hide this option. Used externally within jobs to 

43 determine whether to define job arguments. This is typically used 

44 when you want to track file I/O used by a job but isn't directly 

45 specified in their commands. Off by default. 

46 

47 Examples 

48 -------- 

49 >>> Argument("command", "run").vars() 

50 'run' 

51 

52 >>> files = ["input_1.txt", "input_2.txt"] 

53 >>> Argument("input-files", files).vars() 

54 'input_1.txt input_2.txt' 

55 

56 """ 

57 

58 name: str 

59 argument: int | float | str | list 

60 track: bool = True 

61 suppress: bool = False 

62 static: ClassVar[bool] = False 

63 _args: list[str] = field(init=False) 

64 

65 def __post_init__(self) -> None: 

66 # check against list of protected condor names/characters, 

67 # rename condor variables name to avoid issues 

68 self.condor_name = self.name.replace("-", "_") 

69 if self.condor_name in PROTECTED_CONDOR_VARS: 

70 self.condor_name += "_" 

71 

72 if isinstance(self.argument, str) or not isinstance(self.argument, Iterable): 

73 self.argument = [self.argument] 

74 self._args = [str(arg) for arg in self.argument] 

75 

76 def args(self) -> list[str]: 

77 return self._args 

78 

79 def vars(self, *, basename: bool | Callable[[str], bool] = False) -> str: 

80 args = [] 

81 for arg in self._args: 

82 args.append(path.normalize(arg, basename=basename)) 

83 return " ".join(args) 

84 

85 def files(self, *, basename: bool | Callable[[str], bool] = False) -> str: 

86 args = [] 

87 for arg in self._args: 

88 args.append(path.normalize(arg, basename=basename)) 

89 return ",".join(args) 

90 

91 def remaps(self) -> str: 

92 args = [] 

93 for arg in self._args: 

94 if (normalized := path.normalize(arg, basename=path.is_abs_or_url)) != arg: 

95 args.append(f"{normalized}={arg}") 

96 return ";".join(args) 

97 

98 

99@dataclass 

100class Option: 

101 """Defines a command-line option (long form). 

102 

103 This provides some extra functionality over defining command line 

104 options explicitly, in addition to some extra parameters which 

105 sets how condor interprets how to handle them within the DAG 

106 and within submit descriptions. 

107 

108 Parameters 

109 ---------- 

110 name 

111 The option name to be used in a command. 

112 argument 

113 The argument value(s) used in a command. 

114 track 

115 Whether to track files defined here and used externally within 

116 jobs to determine parent-child relationships when nodes specify 

117 this option as an input or output. On by default. 

118 suppress 

119 Whether to hide this option. Used externally within jobs to 

120 determine whether to define job arguments. This is typically used 

121 when you want to track file I/O used by a job but isn't directly 

122 specified in their commands. Off by default. 

123 prefix 

124 The option prefix to use, e.g. for --verbose, -- is the prefix. 

125 Uses -- by default. 

126 

127 Examples 

128 -------- 

129 >>> Option("verbose").vars() 

130 '--verbose' 

131 

132 >>> Option("input-type", "file").vars() 

133 '--input-type file' 

134 

135 >>> Option("ifos", ["H1", "L1", "V1"]).vars() 

136 '--ifos H1 --ifos L1 --ifos V1' 

137 

138 """ 

139 

140 name: str 

141 argument: int | float | str | list | None = None 

142 track: bool | None = True 

143 suppress: bool = False 

144 prefix: str = "--" 

145 static: ClassVar[bool] = False 

146 _args: list[str] = field(init=False) 

147 

148 def __post_init__(self) -> None: 

149 # check against list of protected condor names/characters, 

150 # rename condor variables name to avoid issues 

151 self.condor_name = self.name.replace("-", "_") 

152 if self.condor_name in PROTECTED_CONDOR_VARS: 

153 self.condor_name += "_" 

154 

155 if self.argument is not None: 

156 if isinstance(self.argument, str) or not isinstance( 

157 self.argument, Iterable 

158 ): 

159 self.argument = [self.argument] 

160 self._args = [str(arg) for arg in self.argument] 

161 

162 def args(self) -> list[str]: 

163 return self._args 

164 

165 def vars(self, *, basename: bool | Callable[[str], bool] = False) -> str: 

166 if self.argument is None: 

167 return f"{self.prefix}{self.name}" 

168 args = [] 

169 for arg in self._args: 

170 normalized_path = path.normalize(arg, basename=basename) 

171 args.append(f"{self.prefix}{self.name} {normalized_path}") 

172 return " ".join(args) 

173 

174 def files(self, *, basename: bool | Callable[[str], bool] = False) -> str: 

175 args = [] 

176 for arg in self._args: 

177 args.append(path.normalize(arg, basename=basename)) 

178 return ",".join(args) 

179 

180 def remaps(self) -> str: 

181 args = [] 

182 for arg in self._args: 

183 if (normalized := path.normalize(arg, basename=path.is_abs_or_url)) != arg: 

184 args.append(f"{normalized}={arg}") 

185 return ";".join(args) 

186 

187 

188@dataclass 

189class Literal: 

190 """Defines a command-line literal. 

191 

192 This provides some extra functionality over defining command line 

193 argument explicitly, in addition to some extra parameters which 

194 sets how condor interprets how to handle them within the DAG 

195 and within submit descriptions. 

196 

197 Parameters 

198 ---------- 

199 argument 

200 The positional argument value(s) used in a command. 

201 track 

202 Whether to track files defined here and used externally within 

203 jobs to determine parent-child relationships when nodes specify 

204 this option as an input or output. On by default. 

205 

206 Examples 

207 -------- 

208 >>> Literal("run").vars() 

209 'run' 

210 

211 >>> Literal("/path/to/input_1.txt").remaps() 

212 'input_1.txt=/path/to/input_1.txt' 

213 

214 """ 

215 

216 argument: int | float | str 

217 track: bool = True 

218 static: ClassVar[bool] = True 

219 

220 @property 

221 def name(self) -> str: 

222 return str(self.argument) 

223 

224 def args(self) -> list[str]: 

225 return [self.name] 

226 

227 def vars(self, *, basename: bool | Callable[[str], bool] = False) -> str: 

228 return path.normalize(self.name, basename=basename) 

229 

230 def files(self, *, basename: bool | Callable[[str], bool] = False) -> str: 

231 return path.normalize(self.name, basename=basename) 

232 

233 def remaps(self) -> str: 

234 normalized = path.normalize(self.name, basename=path.is_abs_or_url) 

235 if normalized != self.name: 

236 return f"{normalized}={self.name}" 

237 return ""