Python ast.get_docstring() Examples

def get_docstring_and_version_via_ast(target):
    """
    Return a tuple like (docstring, version) for the given module,
    extracted by parsing its AST.
    """
    # read as bytes to enable custom encodings
    with target.file.open('rb') as f:
        node = ast.parse(f.read())
    for child in node.body:
        # Only use the version from the given module if it's a simple
        # string assignment to __version__
        is_version_str = (isinstance(child, ast.Assign) and
                          len(child.targets) == 1 and
                          child.targets[0].id == "__version__" and
                          isinstance(child.value, ast.Str))
        if is_version_str:
            version = child.value.s
            break
    else:
        version = None
    return ast.get_docstring(node), version 

def get_func_comments(function_definitions):
    doc = ''
    for f in function_definitions:
        temp_str = to_md(parse_func_string(ast.get_docstring(f)))
        doc += ''.join(
            [
                '### ',
                f.name.replace('_', '\\_'),
                '\n',
                '```python',
                '\n',
                'def ',
                f.name,
                parse_func_args(f),
                '\n',
                '```',
                '\n',
                temp_str,
                '\n',
            ]
        )

    return doc 

def get_comments_str(file_name):
    with open(file_name) as fd:
        file_contents = fd.read()
    module = ast.parse(file_contents)

    function_definitions = [node for node in module.body if isinstance(node, ast.FunctionDef)]

    doc = get_func_comments(function_definitions)

    class_definitions = [node for node in module.body if isinstance(node, ast.ClassDef)]
    for class_def in class_definitions:
        temp_str = to_md(parse_func_string(ast.get_docstring(class_def)))

        # excludes private methods (start with '_')
        method_definitions = [
            node
            for node in class_def.body
            if isinstance(node, ast.FunctionDef) and (node.name[0] != '_' or node.name[:2] == '__')
        ]

        temp_str += get_func_comments(method_definitions)
        doc += '## class ' + class_def.name + '\n' + temp_str
    return doc 

def get_func_comments(function_definitions):
    doc = ''
    for f in function_definitions:
        temp_str = to_md(parse_func_string(ast.get_docstring(f)))
        doc += ''.join(
            [
                '### ',
                f.name.replace('_', '\\_'),
                '\n',
                '```python',
                '\n',
                'def ',
                f.name,
                parse_func_args(f),
                '\n',
                '```',
                '\n',
                temp_str,
                '\n',
            ]
        )

    return doc 

def get_actor_io_ids(filename):
    '''
    Parse IMAS Python actor script and return actor input and output IDSs

    :param filename: filename of the IMAS Python actor

    :return: tuple with list of input IDSs and output IDSs
    '''
    import ast
    with open(filename, 'r') as f:
        module = ast.parse(f.read())
    actor = os.path.splitext(os.path.split(filename)[-1])[0]
    function_definitions = [node for node in module.body if isinstance(node, ast.FunctionDef)]
    docstring = ast.get_docstring([f for f in function_definitions if f.name == actor][0])
    ids_in = []
    ids_out = []
    for line in docstring.split('\n'):
        if 'codeparams' in line:
            pass
        elif line.strip().startswith(':param result:'):
            ids_out = list(map(lambda x: x.strip()[:-1], line.split(':')[2].strip(', ').split(',')))
            break
        elif line.strip().startswith(':param '):
            ids_in.append(line.split(':')[2].strip())
    return ids_in, ids_out 

def test_multiple_line_item_definition_cyk(self):
        """Make sure item definitions can span multiple lines."""
        func = '\n'.join([
            'def do_nothing(x):',
            '    """Do nothing with x.',
            '    ',
            '    :param x: This is an argument which must be ',
            '        qualified by a large amount of text.',
            '',
            '    """',
            '    pass',
        ])
        doc = ast.get_docstring(ast.parse(func).body[0])
        node = parse(condense(lex(doc)))
        self.assertTrue(
            CykNodeUtils.contains(node, 'arguments-section'),
            node,
        ) 

def visit_ClassDef(self, tree):
        bases = (T(', ').join(map(self.visit, tree.bases)) +
                 ','*(len(tree.bases) == 1))
        decoration = T('{}')
        for decorator in tree.decorator_list:
            decoration = decoration.format(T('{}({})').format(self.visit(decorator), T('{}')))
        ns = self.next_child()
        body = ns.many_to_one(tree.body, after=T('{__l}'))
        doc = ast.get_docstring(tree, clean=False)
        body = self.close(ns, "{{'__module__': __name__{}}}".format(
            '' if doc is None else ", '__doc__': {!r}".format(doc)), body)
        if tree.bases:
            class_code = T("(lambda b, d: d.get('__metaclass__', getattr(b[0], "
                           "'__class__', type(b[0])))({!r}, b, d))(({}), "
                           "{})").format(tree.name, bases, body)
        else:
            class_code = T("(lambda d: d.get('__metaclass__', {__g}.get("
                           "'__metaclass__', {__types}.ClassType))({!r}, (), "
                           "d))({})").format(tree.name, body)
        class_code = decoration.format(class_code)
        return assignment_component(T('{after}'), self.store_var(tree.name), class_code) 

def test_parser_sections_correctly(self):
        program = '\n'.join([
            'def func(x, l):',
            '    """Add an item to the head of the list.',
            '    ',
            '    :param x: The item to add to the list.',
            '    :return: The list with the item attached.',
            '    ',
            '    """',
            '    return l.appendleft(x)',
        ])
        doc = ast.get_docstring(ast.parse(program).body[0])
        tokens = condense(lex(doc))
        node = parse(tokens)
        self.assertTrue(
            CykNodeUtils.contains(node, 'returns-section'),
        )
        self.assertTrue(
            CykNodeUtils.contains(node, 'arguments-section'),
        ) 

def get_comments_str(file_name):
    with open(file_name) as fd:
        file_contents = fd.read()
    module = ast.parse(file_contents)

    function_definitions = [node for node in module.body if isinstance(node, ast.FunctionDef)]

    doc = get_func_comments(function_definitions)

    class_definitions = [node for node in module.body if isinstance(node, ast.ClassDef)]
    for class_def in class_definitions:
        temp_str = to_md(parse_func_string(ast.get_docstring(class_def)))

        # excludes private methods (start with '_')
        method_definitions = [
            node
            for node in class_def.body
            if isinstance(node, ast.FunctionDef) and (node.name[0] != '_' or node.name[:2] == '__')
        ]

        temp_str += get_func_comments(method_definitions)
        doc += '## class ' + class_def.name + '\n' + temp_str
    return doc 

def test_parse_noqa_for_argument(self):
        func = '\n'.join([
            'def my_function():',
            '    """Has an extra argument, but thats okay.',
            '',
            '    Args:',
            '        arg1: This will be defined very soon.  # noqa: I102',
            '',
            '    """',
            '    print("Not done yet!")',
        ])
        doc = ast.get_docstring(ast.parse(func).body[0])
        self.assertTrue(doc.startswith('Has an extra'))
        node = parse(condense(lex(doc)))
        self.assertTrue(
            CykNodeUtils.contains(node, 'noqa'),
        )
        noqa = self.get_identifier(node, NoqaIdentifier)
        self.assertTrue(
            noqa is not None,
        )
        self.assertEqual(
            NoqaIdentifier.extract(noqa),
            'I102',
        ) 

def parse_class(node):
    """
    parse_class parses the given node representing a test class for example
    test cases and puts everything into a Section object.
    """
    name = node.name[4:]
    title, details = parse_docstring(ast.get_docstring(node))
    examples = []

    for n in node.body:
        if isinstance(n, _ast.FunctionDef) and n.name.startswith('test_'):
            example = parse_function(n)
            examples.append(example._replace(
                name='{}__{}'.format(name, example.name)))

    return Section(name, title, details, examples) 

def test_parse_long_description_multiple_sections(self):
        func = '\n'.join([
            'def this_function_has_multiple_long_descriptions():',
            '    """Do some math.',
            '',
            '    This is the first part of the long description.',
            '    it can be multiple lines, but doesn\'t have to be',
            '',
            '    This is the second half of the long description.',
            '',
            '    And the final part of it.',
            '',
            '    """',
            '    pass',
        ])
        doc = ast.get_docstring(ast.parse(func).body[0])
        tokens = list(lex(doc))
        node = parse(tokens)
        self.assertTrue(node is not None) 

def test_parses_long_description(self):
        func = '\n'.join([
            'def this_function_has_a_long_description(arg1):',
            '    """Return the arg, unchanged.',
            '',
            '    This function returns the arg, unchanged.  There is',
            '    no particular reason, but this is a good place to check to ',
            '    see that long descriptions are being parsed correctly. ',
            '    If they are, I\'m not sure why.  There is some magic ',
            '    going on here, in fact.',
            '',
            '    """',
            '    return arg1',
        ])
        doc = ast.get_docstring(ast.parse(func).body[0])
        tokens = list(lex(doc))
        node = parse(tokens)
        self.assertTrue(node is not None) 

def parse_function(node):
    old_style = None
    new_style = None
    output = None
    setup = []
    setup_done = False
    title, details = parse_docstring(ast.get_docstring(node, clean=True))
    name = node.name[5:] if node.name.startswith('test_') else node.name

    for n in node.body:
        # Ignore the docstring
        if isinstance(n, _ast.Expr) and isinstance(n.value, _ast.Str):
            continue
        if isinstance(n, _ast.Assign) and n.targets[0].id == 'old_result':
            setup_done = True
            old_style = unparse(n.value, strip=True)
        if isinstance(n, _ast.Assign) and n.targets[0].id == 'new_result':
            setup_done = True
            new_style = unparse(n.value, strip=True)
        if isinstance(n, _ast.Assert) and isinstance(
                n.test.comparators[0], _ast.Str):
            setup_done = True
            output = n.test.comparators[0].s
        if not setup_done:
            setup.append(n)

    if setup:
        setup = unparse(setup, strip=True)

    return Example(
        name,
        title,
        details,
        setup or "",
        old_style or "",
        new_style or "",
        output or ""
    ) 

def _mark_docstring_sub_nodes(node):
        """
        Inspired by ast.get_docstring, mark all docstring sub nodes.

        Case1:
        regular docstring of function/class/module

        Case2:
        def foo(self):
            '''pure string expression'''
            for x in self.contents:
                '''pure string expression'''
                print x
            if self.abc:
                '''pure string expression'''
                pass

        Case3:
        def foo(self):
            if self.abc:
                print('ok')
            else:
                '''pure string expression'''
                pass

        :param node: every ast node
        :return:
        """

        def _mark_docstring_nodes(body):
            if body and isinstance(body, collections.Sequence):
                for n in body:
                    if isinstance(n, ast.Expr) and isinstance(n.value, ast.Str):
                        n.is_docstring = True

        node_body = getattr(node, 'body', None)
        _mark_docstring_nodes(node_body)
        node_orelse = getattr(node, 'orelse', None)
        _mark_docstring_nodes(node_orelse) 

def _parse_docstring(node):
        """Extract code from docstring."""
        docstring = ast.get_docstring(node)
        if docstring:
            parser = doctest.DocTestParser()
            try:
                dt = parser.get_doctest(docstring, {}, None, None, None)
            except ValueError:
                # >>> 'abc'
                pass
            else:
                examples = dt.examples
                return '\n'.join([example.source for example in examples])
        return None 

def visit_Module(self, node: Module):
        """Upon visiting a module, initialize :attr:`DocStringCoverageVisitor.tree`
        with module-wide node info."""
        has_doc = get_docstring(node) is not None and get_docstring(node).strip() != ""
        is_empty = not len(node.body)
        self.tree.append((has_doc, is_empty, []))
        self.generic_visit(node) 

def load(cls, root, klass, name) -> Optional["ModuleInfo"]:
        names = ['SCHEMA']

        def value(a):
            v = a.value
            if isinstance(v, ast.Str):
                return v.s
            return ""

        def filter_type(lst, target):
            return [x for x in lst if isinstance(x, target)]

        def targets(a):
            return [t.id for t in filter_type(a.targets, ast.Name)]

        base = cls.module_class_to_directory(klass)
        if not base:
            raise ValueError(f"Unsupported type: {klass}")

        path = os.path.join(root, base, name)
        try:
            with open(path) as f:
                data = f.read()
        except FileNotFoundError:
            return None

        tree = ast.parse(data, name)

        docstring = ast.get_docstring(tree)
        doclist = docstring.split("\n")

        assigns = filter_type(tree.body, ast.Assign)
        targets = [(t, a) for a in assigns for t in targets(a)]
        values = {k: value(v) for k, v in targets if k in names}
        info = {
            'schema': values.get("SCHEMA"),
            'desc': doclist[0],
            'info': "\n".join(doclist[1:])
        }
        return cls(klass, name, info) 

def test_get_docstring(self):
        node = ast.parse('def foo():\n  """line one\n  line two"""')
        self.assertEqual(ast.get_docstring(node.body[0]),
                         'line one\nline two') 

def class_check(node, kw=False):
    '''
    Class specific check
        Action - check private method, move to utils
        check docstring
        scan for class - recursive class check
        scan for function - function check
    '''
    status = True

    for child in ast.iter_child_nodes(node):
        if isinstance(child, ast.FunctionDef):
            if kw and child.name.startswith("_") and child.name != "__init__":
                print node.name, child.name, "should move to utils"
                status = False
            tmp_status = func_check(child, kw)
            status &= tmp_status
        elif isinstance(child, ast.ClassDef):
            tmp_status = class_check(child, kw)
            status &= tmp_status

    if ast.get_docstring(node) is None:
        # check for docstring
        print node.name, "doesn't contain any docstring"
        status = False
    return status 

def _mark_docstring_sub_nodes(node):
        """
        Inspired by ast.get_docstring, mark all docstring sub nodes.

        Case1:
        regular docstring of function/class/module

        Case2:
        def foo(self):
            '''pure string expression'''
            for x in self.contents:
                '''pure string expression'''
                print x
            if self.abc:
                '''pure string expression'''
                pass

        Case3:
        def foo(self):
            if self.abc:
                print('ok')
            else:
                '''pure string expression'''
                pass

        :param node: every ast node
        :return:
        """

        def _mark_docstring_nodes(body):
            if body and isinstance(body, collections.Sequence):
                for n in body:
                    if isinstance(n, ast.Expr) and isinstance(n.value, ast.Str):
                        n.is_docstring = True

        node_body = getattr(node, 'body', None)
        _mark_docstring_nodes(node_body)
        node_orelse = getattr(node, 'orelse', None)
        _mark_docstring_nodes(node_orelse) 

def test_get_docstring(self):
        node = ast.parse('def foo():\n  """line one\n  line two"""')
        self.assertEqual(ast.get_docstring(node.body[0]),
                         'line one\nline two')

        node = ast.parse('async def foo():\n  """spam\n  ham"""')
        self.assertEqual(ast.get_docstring(node.body[0]), 'spam\nham') 

def test_get_docstring(self):
        node = ast.parse('def foo():\n  """line one\n  line two"""')
        self.assertEqual(ast.get_docstring(node.body[0]),
                         'line one\nline two') 

def test_get_docstring(self):
        node = ast.parse('def foo():\n  """line one\n  line two"""')
        self.assertEqual(ast.get_docstring(node.body[0]),
                         'line one\nline two') 

def _visit_helper(self, node):
        """Helper method to update :attr:`DocStringCoverageVisitor.tree` with pertinent
        documentation information for `node`, then ensure all child nodes are
        also visited"""
        self.symbol_count += 1
        has_doc = get_docstring(node) is not None and get_docstring(node).strip() != ""
        _node = (node.name, has_doc, [])
        self.tree[-1][-1].append(_node)
        self.tree.append(_node)
        self.generic_visit(node)
        self.tree.pop() 

def __init__(self, function_def, parent_class=None, testmodule=None):
        """Wrap a ``ast.FunctionDef`` instance used to extract information."""
        self.docstring = ast.get_docstring(function_def)
        self.function_def = function_def
        self.name = function_def.name
        if parent_class:
            self.parent_class = parent_class.name
            self.parent_class_def = parent_class
            self.class_docstring = ast.get_docstring(self.parent_class_def)
        else:
            self.parent_class = None
            self.parent_class_def = None
            self.class_docstring = None
        self.testmodule = testmodule.path
        self.module_def = testmodule
        self.module_docstring = ast.get_docstring(self.module_def)
        self.pkginit = os.path.join(
            os.path.dirname(self.testmodule), '__init__.py')
        if os.path.exists(self.pkginit):
            self.pkginit_def = ast.parse(''.join(open(self.pkginit)))
            self.pkginit_docstring = ast.get_docstring(self.pkginit_def)
        else:
            self.pkginit_def = None
            self.pkginit_docstring = None
        self.tokens = {}
        self.invalid_tokens = {}
        self._rst_parser_messages = []
        tokens = SETTINGS.get('tokens').keys() or None
        minimum_tokens = [key for key, value
                          in SETTINGS.get('tokens').items()
                          if value.required] or None
        self.parser = DocstringParser(tokens, minimum_tokens)
        self._parse_docstring()
        self._parse_decorators() 

def get_init_code(tree, table):
    """Get one-lined code from `tree` and `table.

    Calculate the helper variables that we will need, and wrap the output
    code (computed from `tree`) with definitions of those variables.

    TODO: Short-circuit to something far simpler if the program has only one
    print statement.

    Arguments:
        tree: Python AST node
        table: symtable.symtable

    Returns:
        string - valid one-line code
    """
    output = Namespace(table).many_to_one(tree.body)

    doc = ast.get_docstring(tree, clean=False)
    if doc is not None:
        output = assignment_component(output, T("{__g}['__doc__']"), repr(doc))

    output = provide(
        output.format(__l=T('{__g}')),
        __print=T("__import__('__builtin__', level=0).__dict__['print']"),
        __y="(lambda f: (lambda x: x(x))(lambda y:"
          " f(lambda: y(y)())))",
        __g=T("globals()"),
        __contextlib="__import__('contextlib', level=0)",
        __operator="__import__('operator', level=0)",
        __sys="__import__('sys', level=0)",
        __types="__import__('types', level=0)")

    return output.close() 

def visit_FunctionDef(self, tree):
        # self.visit() returns something of the form
        # ('lambda x, y, z=5, *args: ', ['x', 'y', 'z', 'args'])
        args, arg_names = self.visit(tree.args)
        decoration = T('{}')
        for decorator in tree.decorator_list:
            decoration = decoration.format(T('{}({})').format(self.visit(decorator), T('{}')))
        ns = self.next_child()
        body = ns.many_to_one(tree.body).format(pre_return='', post_return='')
        if arg_names:
            body = assignment_component(body,
                T(', ').join(ns.var(name) for name in arg_names),
                T(', ').join(arg_names))
        body = self.close(ns, '{}', body)
        function_code = args + body
        doc = ast.get_docstring(tree, clean=False)
        if tree.decorator_list:
            return assignment_component(
                T('{after}'),
                self.store_var(tree.name),
                decoration.format(assignment_component(
                    '__func',
                    '__func, __func.__name__' + ('' if doc is None else ', __func.__doc__'),
                    T('{}, {!r}' + ('' if doc is None else ', {!r}')).format(
                        function_code, tree.name, doc))))
        else:
            return assignment_component(
                T('{after}'),
                T('{}, {}.__name__' + ('' if doc is None else ', {}.__doc__')).format(
                    self.store_var(tree.name), self.var(tree.name), self.var(tree.name)),
                T('{}, {!r}' + ('' if doc is None else ', {!r}')).format(
                    function_code, tree.name, doc)) 

def test_get_docstring(self):
        node = ast.parse('def foo():\n  """line one\n  line two"""')
        self.assertEqual(ast.get_docstring(node.body[0]),
                         'line one\nline two') 

def _parse_docstring(node):
        """Extract code from docstring."""
        docstring = ast.get_docstring(node)
        if docstring:
            parser = doctest.DocTestParser()
            try:
                dt = parser.get_doctest(docstring, {}, None, None, None)
            except ValueError:
                # >>> 'abc'
                pass
            else:
                examples = dt.examples
                return '\n'.join([example.source for example in examples])
        return None