Coverage for src/pyTRLCConverter/plantuml.py: 66%

1"""PlantUML to image file converter. 

2 

3 Author: Andreas Merkle (andreas.merkle@newtec.de) 

4""" 

5 

6# pyTRLCConverter - A tool to convert TRLC files to specific formats. 

7# Copyright (c) 2024 - 2026 NewTec GmbH 

8# 

9# This file is part of pyTRLCConverter program. 

10# 

11# The pyTRLCConverter program is free software: you can redistribute it and/or modify it under 

12# the terms of the GNU General Public License as published by the Free Software Foundation, 

13# either version 3 of the License, or (at your option) any later version. 

14# 

15# The pyTRLCConverter program is distributed in the hope that it will be useful, but 

16# WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or 

17# FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details. 

18# 

19# You should have received a copy of the GNU General Public License along with pyTRLCConverter. 

20# If not, see <https://www.gnu.org/licenses/>. 

21 

22# Imports ********************************************************************** 

23import os 

24import subprocess 

25import sys 

26import zlib 

27import base64 

28import urllib 

29import urllib.parse 

30import urllib3 

31import requests 

32 

33from pyTRLCConverter.logger import log_verbose, log_error 

34 

35# Variables ******************************************************************** 

36 

37# URL encoding char sets. 

38# See https://plantuml.com/de/text-encoding for differences to base64 in URL encode. 

39BASE64_ENCODE_CHARS = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/" 

40PLANTUML_ENCODE_CHARS = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz-_" 

41PLANTUML_ENV_VAR = "PLANTUML" 

42PLANTUML_VERIFY_SSL_ENV_VAR = "PLANTUML_VERIFY_SSL" 

43 

44# Classes ********************************************************************** 

45 

46 

47class PlantUML(): 

48 # lobster-trace: SwRequirements.sw_req_plantuml 

49 """PlantUML image generator. 

50 """ 

51 def __init__(self) -> None: 

52 self._server_url = None 

53 self._plantuml_jar = None 

54 self._working_directory = os.path.abspath(os.getcwd()) 

55 # Default to True; set PLANTUML_VERIFY_SSL=false to disable for internal servers 

56 # with self-signed or corporate CA certificates. 

57 self._verify_ssl = os.environ.get(PLANTUML_VERIFY_SSL_ENV_VAR, "true").lower() != "false" 

58 if not self._verify_ssl: 

59 urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning) 

60 

61 if PLANTUML_ENV_VAR in os.environ: 

62 plantuml_access = os.environ[PLANTUML_ENV_VAR] 

63 try: 

64 # Use server method if PLANTUML env var points to a URL. 

65 if urllib.parse.urlparse(plantuml_access).scheme in ['http', 'https']: 

66 self._server_url = plantuml_access 

67 else: 

68 # Otherwise assume it's a local JAR. 

69 self._plantuml_jar = plantuml_access 

70 except ValueError: 

71 self._plantuml_jar = plantuml_access 

72 

73 def _get_absolute_path(self, path): 

74 """Convert a relative path to an absolute path based on the working directory. 

75 

76 The plantuml.jar requires an absolute output path. 

77 

78 Args: 

79 path (str): Relative or absolute path. 

80 

81 Returns: 

82 str: Absolute path. 

83 """ 

84 absolute_path = path 

85 

86 if os.path.isabs(path) is False: 

87 absolute_path = os.path.join(self._working_directory, path) 

88 

89 return absolute_path 

90 

91 def is_plantuml_file(self, diagram_path): 

92 """Is the diagram a PlantUML file? 

93 Only the file extension will be checked. 

94 

95 Args: 

96 diagram_path (str): Diagram path. 

97 

98 Returns: 

99 bool: If file is a PlantUML file, it will return True otherwise False. 

100 """ 

101 is_plantuml = False 

102 

103 if diagram_path.endswith(".plantuml") or \ 

104 diagram_path.endswith(".puml") or \ 

105 diagram_path.endswith(".wsd"): 

106 is_plantuml = True 

107 

108 return is_plantuml 

109 

110 def _make_server_url(self, diagram_type: str, diagram_source: str, source_is_file: bool = True) -> str: 

111 """Generate a PlantUML server GET URL from a diagram data file. 

112 

113 Args: 

114 diagram_type (str): Diagram type, e.g. svg. See PlantUML -t options. 

115 diagram_source (str): Path to the PlantUML diagram or content of PlantUML diagram. 

116 source_is_file (bool): True if diagram_source points to a file, False if it is the diagram content. 

117 

118 Raises: 

119 FileNotFoundError: PlantUML diagram file not found. 

120 """ 

121 if source_is_file: 

122 # Read PlantUML diagram data from given file. 

123 with open(diagram_source, 'r', encoding='utf-8') as input_file: 

124 diagram_string = input_file.read().encode('utf-8') 

125 else: 

126 # Use the provided source directly as the diagram_string. 

127 diagram_string = diagram_source.encode('utf-8') 

128 

129 # Compress the data using deflate. 

130 # Strip Zlib's 2 byte header and 4 byte checksum for raw deflate data. 

131 compressed_data = zlib.compress(diagram_string)[2:-4] 

132 

133 # Encode the compressed data using base64. 

134 base64_encoded_data = base64.b64encode(compressed_data) 

135 

136 # Translate from base64 to plantuml char encoding. 

137 base64_to_puml_trans = bytes.maketrans( 

138 BASE64_ENCODE_CHARS.encode('utf-8'), 

139 PLANTUML_ENCODE_CHARS.encode('utf-8') 

140 ) 

141 puml_encoded_data = base64_encoded_data.translate(base64_to_puml_trans).decode('utf-8') 

142 

143 # Create the URL for the PlantUML server. 

144 query_url = ( 

145 f"{self._server_url}/" 

146 f"{diagram_type}/" 

147 f"{urllib.parse.quote(puml_encoded_data)}" 

148 ) 

149 

150 return query_url 

151 

152 def generate_to_bytes(self, diagram_type: str, diagram_source: str) -> bytes: 

153 """Generate a PlantUML image from source text and return the raw bytes. 

154 

155 Args: 

156 diagram_type (str): Diagram type, e.g. png. See PlantUML -t options. 

157 diagram_source (str): PlantUML diagram source text. 

158 

159 Raises: 

160 FileNotFoundError: PlantUML jar file not found or Java not installed (local mode). 

161 requests.exceptions.RequestException: HTTP error from PlantUML server (server mode). 

162 

163 Returns: 

164 bytes: The raw image bytes. 

165 """ 

166 assert diagram_type in ("png", "svg") 

167 

168 if self._server_url is not None: 

169 result = self._generate_to_bytes_server(diagram_type, diagram_source) 

170 else: 

171 result = self._generate_to_bytes_local(diagram_type, diagram_source) 

172 

173 return result 

174 

175 def _generate_to_bytes_server(self, diagram_type: str, diagram_source: str) -> bytes: 

176 """Generate image via PlantUML server and return raw bytes. 

177 

178 Args: 

179 diagram_type (str): Diagram type, e.g. png. 

180 diagram_source (str): PlantUML diagram source text. 

181 

182 Raises: 

183 requests.exceptions.RequestException: HTTP error from PlantUML server. 

184 

185 Returns: 

186 bytes: The raw image bytes. 

187 """ 

188 

189 url = self._make_server_url(diagram_type, diagram_source, source_is_file=False) 

190 log_verbose(f"Sending GET request {url}") 

191 response = requests.get(url, timeout=10, verify=self._verify_ssl) 

192 

193 if response.status_code != 200: 

194 raise requests.exceptions.RequestException( 

195 f"{response.status_code} - {response.text}" 

196 ) 

197 

198 return response.content 

199 

200 def _generate_to_bytes_local(self, diagram_type: str, diagram_source: str) -> bytes: 

201 """Generate image via local plantuml.jar using -pipe mode and return raw bytes. 

202 

203 Args: 

204 diagram_type (str): Diagram type, e.g. png. 

205 diagram_source (str): PlantUML diagram source text. 

206 

207 Raises: 

208 FileNotFoundError: Java not installed or plantuml.jar not found. 

209 

210 Returns: 

211 bytes: The raw image bytes. 

212 """ 

213 if self._plantuml_jar is None: 

214 raise FileNotFoundError( 

215 f"PlantUML not found. Set the {PLANTUML_ENV_VAR} environment variable" 

216 " to the path of plantuml.jar or a server URL." 

217 ) 

218 

219 if not os.path.isfile(self._plantuml_jar): 

220 raise FileNotFoundError(f"plantuml.jar at {self._plantuml_jar} not found.") 

221 

222 plantuml_cmd = ["java"] 

223 

224 if sys.platform.startswith("linux"): 

225 plantuml_cmd.append("-Djava.awt.headless=true") 

226 

227 plantuml_cmd.extend([ 

228 "-jar", self._plantuml_jar, 

229 f"-t{diagram_type}", 

230 "-pipe" 

231 ]) 

232 

233 try: 

234 output = subprocess.run( 

235 plantuml_cmd, 

236 input=diagram_source.encode("utf-8"), 

237 capture_output=True, 

238 check=False 

239 ) 

240 except FileNotFoundError as exc: 

241 raise FileNotFoundError( 

242 "Java not found. Ensure Java is installed and available on PATH." 

243 ) from exc 

244 

245 if output.returncode != 0: 

246 output.check_returncode() 

247 

248 return output.stdout 

249 

250 def generate(self, diagram_type: str, diagram_path: str, dst_path: str) -> None: 

251 """Generate plantuml image. 

252 

253 Args: 

254 diagram_type (str): Diagram type, e.g. svg. See PlantUML -t options. 

255 diagram_path (str): Path to the PlantUML diagram. 

256 dst_path (str): Path to the destination of the generated image. 

257 

258 Raises: 

259 FileNotFoundError: PlantUML java jar file not found in local mode. 

260 FileNotFoundError: PlantUML diagram file not found. 

261 requests.exceptions.RequestException: Error during GET request to PlantUML server. 

262 OSError: Destination path does not exist. 

263 """ 

264 assert diagram_type in ("png", "svg") 

265 

266 if self._server_url is not None: 

267 self._generate_server(diagram_type, diagram_path, dst_path) 

268 else: 

269 self._generate_local(diagram_type, diagram_path, dst_path) 

270 

271 def _generate_server(self, diagram_type: str, diagram_path: str, dst_path: str) -> None: 

272 """Generate image using a plantuml server. 

273 

274 Does not require Java installed and is usually a lot faster 

275 as no Java startup time is needed when using the plantuml.jar file. 

276 

277 Args: 

278 diagram_type (str): Diagram type, e.g. svg. See PlantUML -t options. 

279 diagram_path (str): Path to the PlantUML diagram. 

280 dst_path (str): Path to the destination of the generated image. 

281 

282 Raises: 

283 FileNotFoundError: PlantUML diagram file not found. 

284 OSError: Destination path does not exist. 

285 requests.exceptions.RequestException: Error during GET request to PlantUML server. 

286 """ 

287 if not os.path.exists(dst_path): 

288 os.makedirs(dst_path) 

289 

290 # Send GET request to the PlantUML server. 

291 url = self._make_server_url(diagram_type, diagram_path) 

292 log_verbose(f"Sending GET request {url}") 

293 response = requests.get(url, timeout=10, verify=self._verify_ssl) 

294 

295 if response.status_code == 200: 

296 # Save the response content in image file. 

297 output_file = os.path.splitext(os.path.basename(diagram_path))[0] 

298 output_file += "." + diagram_type 

299 output_file = os.path.join(dst_path, output_file) 

300 with open(output_file, 'wb') as image_file: 

301 image_file.write(response.content) 

302 

303 log_verbose(f"Diagram saved as {output_file}.") 

304 else: 

305 raise requests.exceptions.RequestException(f"{response.status_code} - {response.text}") 

306 

307 def _generate_local(self, diagram_type, diagram_path, dst_path): 

308 """Generate image local call to plantuml.jar. 

309 

310 Args: 

311 diagram_type (str): Diagram type, e.g. svg. See PlantUML -t options. 

312 diagram_path (str): Path to the PlantUML diagram. 

313 dst_path (str): Path to the destination of the generated image. 

314 

315 Raises: 

316 FileNotFoundError: PlantUML java jar file not found. 

317 OSError: Destination path does not exist. 

318 """ 

319 if self._plantuml_jar is not None: 

320 

321 if not os.path.exists(dst_path): 

322 os.makedirs(dst_path) 

323 

324 plantuml_cmd = ["java"] 

325 

326 if sys.platform.startswith("linux"): 

327 plantuml_cmd.append("-Djava.awt.headless=true") 

328 

329 plantuml_cmd.extend( 

330 [ 

331 "-jar", f"{self._plantuml_jar}", 

332 f"{diagram_path}", 

333 f"-t{diagram_type}", 

334 "-o", self._get_absolute_path(dst_path) 

335 ] 

336 ) 

337 

338 try: 

339 output = subprocess.run(plantuml_cmd, capture_output=True, text=True, check=False) 

340 except FileNotFoundError as exc: 

341 raise FileNotFoundError( 

342 "Java not found. Ensure Java is installed and available on PATH." 

343 ) from exc 

344 

345 if output.stderr: 

346 log_error(output.stderr, True) 

347 if output.returncode != 0: 

348 if not os.path.isfile(self._plantuml_jar): 

349 raise FileNotFoundError(f"plantuml.jar at {self._plantuml_jar} not found.") 

350 if not os.path.isfile(diagram_path): 

351 raise FileNotFoundError(f"Diagram at {diagram_path} not found.") 

352 output.check_returncode() 

353 log_verbose(output.stdout) 

354 else: 

355 raise FileNotFoundError( 

356 f"PlantUML not found. Set the {PLANTUML_ENV_VAR} environment variable" 

357 " to the path of plantuml.jar or a server URL." 

358 ) 

359 

360# Functions ******************************************************************** 

361 

362# Main *************************************************************************