nl2ml/parser/auto_features/get_func_docstrings.py

import os
import shutil
import tempfile
from pyls.workspace import Document, Workspace
from pyls.plugins.references import pyls_references
from pyls.plugins.signature import pyls_signature_help

import ast
import pandas as pd
import inspect
import ast
import sys
import numpy as np


def get_func_docstrings(notebook_cells):
    cells = []
    lengths = []
    for cell in notebook_cells.code_block:
        cell = cell.strip('`').strip()
        if cell.find("%matplotlib inline") > -1:
            cell = cell.replace("%matplotlib inline", "")
        code = cell
        for line in cell.split("\n"):
            if "import" in line:
                try:
                    exec(line)
                except Exception as e:
                    print("Could not run:\n", line, "\n", e)
            if len(line.strip()) > 0 and line.strip()[0] == "!":
                code = code.replace(line, "")

        cells.append(code)
        lengths.append(len(code.split("\n")))

    lengths = np.cumsum(lengths)
    cb_ids = list(notebook_cells.code_block_id)

    cells = "\n".join(cells)
    cells_lines = cells.split("\n")

    tmp = tempfile.mkdtemp()
    workspace = Workspace(tmp, None)

    def create_file(name, content):
        fn = os.path.join(tmp, name)
        with open(fn, 'w') as f:
            f.write(content)
        workspace.put_document('file://' + fn, content, None)

    create_file('code.py', cells)

    DOC_URI = 'file://' + os.path.join(workspace.root_uri, 'code.py')
    doc = Document(DOC_URI, workspace)

    result = []  # {codeblock_id, f_name, docstring}
    parsed = ast.parse(cells)

    for node in ast.walk(parsed):
        if isinstance(node, ast.Call):
            line_call = cells_lines[node.lineno - 1][node.col_offset:node.end_col_offset]
            idx = line_call.find("(")

            position_ref = {'line': node.lineno - 1, 'character': idx - 1 + node.col_offset}
            position_sig = {'line': node.lineno - 1, 'character': idx + 1 + node.col_offset}
            name = line_call[:idx]
            docstr = ""
            try:
                docstr = eval('inspect.getdoc(' + name + ')')
            except Exception:
                # find source through pyls
                signatures = pyls_signature_help(doc, position_sig)['signatures']
                if len(signatures) > 0 and 'documentation' in signatures[0] and len(signatures[0]['documentation']) > 0:
                    docstr = signatures[0]['documentation']
                else:
                    try:
                        uri = pyls_references(doc, position_ref)[0]['uri']
                        source = []
                        if uri.find('site-packages') >= 0:
                            uri_split = uri.split('/')
                            source = uri_split[uri_split.index('site-packages') + 1:]
                            source[-1] = source[-1].split(".")[0]
                        f_name = name.split(".")[-1]
                        try:
                            exec("from " + ".".join(source) + " import " + f_name)
                        except Exception as e:
                            pass
                        try:
                            exec("from " + ".".join(source[:-1]) + " import " + f_name)
                        except Exception as e:
                            pass
                        try:
                            docstr = eval('inspect.getdoc(' + f_name + ')')
                        except Exception as e:
                            if source[0] == "pandas":
                                try:
                                    docstr = eval('inspect.getdoc(' + "pandas.DataFrame." + f_name + ')')
                                except Exception as e:
                                    pass
                            else:
                                pass
                    except Exception as e:
                        print("error in pyls_references, function name", name)
                        pass

            result.append({"code_block_id": cb_ids[np.where(lengths >= node.lineno - 1)[0][0]], "f_name": name,
                           "docstring": docstr})
    shutil.rmtree(tmp)

    return result