Use Python to Parse Python Files with Abstract Syntax Tree

README.md

Parsing a Python File using Python and the Abstract Syntax Tree

This comes in handy often. I want to do something based on a Python file. Maybe I want to pull all the docstrings out and make a markdown file. Maybe I want generate a script that calls all the functions in a python file. Whatever it is, the abstract syntax tree can help. Here is an example for the first use case: I pull out all the docstrings. We used this to create the documentation for hackathon at Clear Technologies a few years ago.

extract_docstring.py

import ast
import os

view_path = "./folder_with_python_files_in_it"
docs_dir = "doc"
docs_path = os.path.join(view_path, docs_dir)
undocumented_count = 0

if not os.path.exists(docs_path):
    os.mkdir(docs_path)

only_files = [f for f in os.listdir(view_path) if os.path.isfile(os.path.join(view_path, f))]

new_md_file = open('documentation.md', "w+")

md_content = '<style>' \
             '.body {margin:14px; font-family: verdana, san-serif; font-size: 18px;}\n' \
             'h1, h2, h3, h4, h5, p, li {font-family: verdana, sans-serif;}\n' \
             '.verb {width:64px; height:22px; text-align: center; padding-top: 4px; padding-bottom: 2px; ' \
             'font-family: verdana, sans-serif; font-weight: bold;}\n' \
             '.get {color: white; background-color: green;}\n' \
             '.post {color: white; background-color: orange;}\n' \
             '.put {color:white; background-color: blue;}\n' \
             '.delete {color: white; background-color: red;}\n' \
             '</style>'
md_content += '![logo](https://vsi5.visualstorageintelligence.com//static/media/logo.17a014dd.svg)\n'
md_content += '# API Documentation\n---\n\n'
md_content += "### <a id='toc' /> Table of Contents:\n"
md_content += "- [Business Unit Views](#business_units)\n"
md_content += "- [Admin Views](#admin_views)\n"
for file in only_files:
    if "_views" in file:

        with open(os.path.join(view_path, file)) as fd:
            file_contents = fd.read()
            module = ast.parse(file_contents)
            function_defs = [node for node in module.body if isinstance(node, ast.FunctionDef)]
            anchor = fd.name.split('/')[2][:-3]
            md_content += "## <a id='" + anchor + "' />" + fd.name.split('/')[2][:-3].replace("_", " ").title() + "\r\n"
            for f in function_defs:
                docstring = ast.get_docstring(f)
                if docstring is not None and len(f.decorator_list) != 0:
                    for i in range(len(f.decorator_list)):
                        for dec_str in f.decorator_list[i].args:
                            section = ast.dump(dec_str).split("\'")[1].replace("<", "").replace(">", "")
                            md_content += "### " + section + "\n"

                        for keyword in f.decorator_list[i].keywords:
                            verb = ast.dump(keyword).split(",")[1].split("\'")[1]
                            md_content += "<div class='verb " + verb.lower() + "'>" + verb + "</div>\n"

                        if len(f.decorator_list[0].keywords) == 0:
                            md_content += "#### Verbs: GET \n"
                    docstring = docstring.replace(":param ", "\n- ")
                    docstring = docstring.replace(":return:", "\n#### Returns: ")
                    md_content += (docstring + "\n [Back to Top](#toc)\n\n---\n")
                elif len(f.decorator_list) != 0:
                    undocumented_count += 1
                    print("Undocumented endpoint: " + f.name)

                if len(f.decorator_list) == 0:
                    print("Function " + f.name + " is not an endpoint.")

new_md_file.write(md_content)
new_md_file.close()

print("You have " + str(undocumented_count) + " undocumented endpoints.")