diff --git a/flask_swagger.py b/flask_swagger.py index d594fdb..dab4d80 100644 --- a/flask_swagger.py +++ b/flask_swagger.py @@ -8,12 +8,16 @@ - Added basic_path parameter """ import inspect +import logging import yaml import re import os from collections import defaultdict +from plantuml_docs import generate_plantuml + +logger = logging.getLogger(__name__) def _sanitize(comment): return comment.replace('\n', '
') if comment else comment @@ -46,28 +50,34 @@ def _doc_from_file(path): return doc -def _parse_docstring(obj, process_doc, from_file_keyword, base_path): - first_line, other_lines, swag = None, None, None - full_doc = inspect.getdoc(obj) - if full_doc: - if from_file_keyword is not None: - from_file = _find_from_file(full_doc, from_file_keyword, base_path) - if from_file: - full_doc_from_file = _doc_from_file(from_file) - if full_doc_from_file: - full_doc = full_doc_from_file - line_feed = full_doc.find('\n') - if line_feed != -1: - first_line = process_doc(full_doc[:line_feed]) - yaml_sep = full_doc[line_feed+1:].find('---') - if yaml_sep != -1: - other_lines = process_doc(full_doc[line_feed+1:line_feed+yaml_sep]) - swag = yaml.full_load(full_doc[line_feed+yaml_sep:]) +def _parse_docstring(obj, process_doc, from_file_keyword, base_path, app): + try: + first_line, other_lines, swag = None, None, None + full_doc = inspect.getdoc(obj) + if full_doc: + if from_file_keyword is not None: + from_file = _find_from_file(full_doc, from_file_keyword, base_path) + if from_file: + full_doc_from_file = _doc_from_file(from_file) + if full_doc_from_file: + full_doc = full_doc_from_file + line_feed = full_doc.find('\n') + if line_feed != -1: + first_line = process_doc(full_doc[:line_feed]) + yaml_sep = full_doc[line_feed+1:].find('---') + if yaml_sep != -1: + other_lines = full_doc[line_feed+1:line_feed+yaml_sep] + other_lines = generate_plantuml(other_lines, app) + other_lines = process_doc(other_lines) + swag = yaml.full_load(full_doc[line_feed+yaml_sep:]) + else: + other_lines = process_doc(full_doc[line_feed+1:]) else: - other_lines = process_doc(full_doc[line_feed+1:]) - else: - first_line = full_doc - return first_line, other_lines, swag + first_line = full_doc + return first_line, other_lines, swag + except Exception as e: + logger.error("Failed to parse docstring for %s: %s", obj, e) + raise def _extract_definitions(alist, level=None): @@ -124,7 +134,8 @@ def _extract_array_defs(source): def swagger(app, prefix=None, process_doc=_sanitize, - from_file_keyword=None, template=None, base_path=""): + from_file_keyword=None, template=None, base_path="", + title="Cool product name", version="0.0.0", description=""): """ Call this from an @app.route method like this @app.route('/spec.json') @@ -146,8 +157,9 @@ def spec(): output = { "swagger": "2.0", "info": { - "version": "0.0.0", - "title": "Cool product name", + "version": version, + "title": title, + "description": generate_plantuml(description, app) } } paths = defaultdict(dict) @@ -183,7 +195,8 @@ def spec(): operations = dict() for verb, method in methods.items(): summary, description, swag = _parse_docstring(method, process_doc, - from_file_keyword, base_path) + from_file_keyword, base_path, + app) if swag is not None: # we only add endpoints with swagger data in the docstrings defs = swag.get('definitions', []) defs = _extract_definitions(defs) diff --git a/plantuml_docs.py b/plantuml_docs.py new file mode 100644 index 0000000..154293d --- /dev/null +++ b/plantuml_docs.py @@ -0,0 +1,73 @@ +import hashlib +import logging +import os +import re + +logger = logging.getLogger(__name__) + +try: + import plantuml +except ImportError: + plantuml = None + +_PLANTUML_RE = re.compile("(@startuml.*?@enduml)", re.MULTILINE|re.DOTALL) +FLASK_SWAGGER_PLANTUML_SERVER = 'FLASK_SWAGGER_PLANTUML_SERVER' +FLASK_SWAGGER_PLANTUML_FOLDER = 'FLASK_SWAGGER_PLANTUML_FOLDER' + +def sub(string, match, replacement): + return string[:match.start()] + replacement + string[match.end():] + +def generate_plantuml(docstring, app): + """ + Generate PlantUML diagrams from the given docstring. + + If the plantuml Python package is not installed, the docstring is returned + unaltered. Otherwise, it performs the following steps: + * Looks for any `@startuml...@enduml` pairs. + * When it finds one, extracts the diagram text and sends it to a PlantUML + server. The default is the public server, but this can be configured by + setting the `FLASK_SWAGGER_PLANTUML_SERVER` in the flask app config. + * The image returned from the server is placed into the application's static + files folder. The location of the static folder is determined from the + `app` object. The subfolder defaults to `uml` but can be configured by + setting the `FLASK_SWAGGER_PLANTUML_FOLDER` in the flask app config. + * The original diagram text is replaced with a markdown link to the generated + image. + """ + if not plantuml: + logger.debug("PlantUML not installed; not generating diagrams") + return docstring + + url=app.config.get(FLASK_SWAGGER_PLANTUML_SERVER, 'http://www.plantuml.com/plantuml/img/') + logger.debug("Using PlantUML server %s", url) + server = plantuml.PlantUML(url=url) + + subfolder = app.config.get(FLASK_SWAGGER_PLANTUML_FOLDER, 'uml') + folder = os.path.join(app.static_folder, subfolder) + if not os.path.exists(folder): + os.mkdir(folder) + logger.debug("Outputting diagrams to %s", folder) + + while True: + match = _PLANTUML_RE.search(docstring) + if not match: + break + uml = match.group(1) + # The same UML data will produce the same filename + filename = hashlib.sha256(uml.encode('utf-8')).hexdigest() + '.png' + output_file = os.path.join(folder, filename) + try: + image_data = server.processes(uml) + with open(output_file, 'wb') as file: + file.write(image_data) + docstring = sub(docstring, match, f'![{filename}]({app.static_url_path}/{subfolder}/{filename})') + except plantuml.PlantUMLConnectionError as e: + docstring = sub(docstring, match, f"Failed to connect to the PlantUML server: {e}") + except plantuml.PlantUMLHTTPError as e: + docstring = sub(docstring, match, f"HTTP error while connection to the PlantUML server: {e}") + except plantuml.PlantUMLError as e: + docstring = sub(docstring, match, f"PlantUML error: {e}") + + return docstring + + diff --git a/setup.py b/setup.py index cef6457..5865868 100644 --- a/setup.py +++ b/setup.py @@ -12,7 +12,7 @@ description='Extract swagger specs from your flask project', author='Atli Thorbjornsson', license='MIT', - py_modules=['flask_swagger', 'build_swagger_spec'], + py_modules=['flask_swagger', 'build_swagger_spec', 'plantuml_docs'], long_description=long_description, install_requires=['Flask>=0.10', 'PyYAML>=5.1'], classifiers=[