Embedded plantuml diagrams in docstrings

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', '<br/>') 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)

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
+
+

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=[