#!/usr/bin/env python3
#
# documentation.py
"""
Decorators to add docstrings to enum members from comments.
"""
#
# Copyright (c) 2020-2021 Dominic Davis-Foster <dominic@davis-foster.co.uk>
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU Lesser General Public License as published by
# the Free Software Foundation; either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
# MA 02110-1301, USA.
#
# stdlib
import ast
import inspect
import re
import sys
import tokenize
import warnings
from enum import Enum, EnumMeta
from textwrap import dedent
from typing import Iterable, Iterator, List, Optional, Sequence, Tuple, TypeVar, Union
# 3rd party
import pygments.token # type: ignore[import]
from pygments.lexers.python import PythonLexer # type: ignore[import]
__all__ = [
"get_tokens",
"document_enum",
"document_member",
"parse_tokens",
"get_base_indent",
"DocumentedEnum",
"get_dedented_line",
"MultipleDocstringsWarning",
]
_lexer = PythonLexer()
INTERACTIVE = bool(getattr(sys, "ps1", sys.flags.interactive))
EnumType = TypeVar("EnumType", bound=EnumMeta)
[docs]def get_tokens(line: str) -> List[Tuple]:
"""
Returns a list ot tokens generated from the given Python code.
:param line: Line of Python code to tokenise.
"""
return list(_lexer.get_tokens(line))
def _docstring_from_expr(expr: ast.Expr) -> Optional[str]:
"""
Check if the expression is a docstring.
:param expr:
:returns: The cleaned docstring text if it is a docstring, or :py:obj:`None` if it isn't.
"""
# might be docstring
docstring_node = expr.value
if isinstance(docstring_node, ast.Constant) and isinstance(docstring_node.value, str):
text = docstring_node.value
elif isinstance(docstring_node, ast.Str):
text = docstring_node.s
else:
# not a docstring
return None
return inspect.cleandoc(text)
def _docstring_from_eol_comment(
source: str,
node: Union[ast.Assign, ast.AnnAssign],
) -> Optional[str]:
"""
Search for an end-of-line docstring comment (starts with ``# doc:``).
:param source: The source of the Enum class.
:param node: The AST node for the Enum member.
"""
toks = _tokenize_line(source.split('\n')[node.lineno - 1])
comment_toks = [x for x in list(toks) if x.type == tokenize.COMMENT]
if comment_toks:
for match in re.finditer(r"(doc:\s*)([^#]*)(#|$)", comment_toks[0].string):
if match.group(2):
return match.group(2).rstrip()
return None
def _docstring_from_sphinx_comment(
source: str,
node: Union[ast.Assign, ast.AnnAssign],
) -> Optional[str]:
"""
Search for a Sphinx-style docstring comment (starts with ``#:``).
:param source: The source of the Enum class.
:param node: The AST node for the Enum member.
"""
for offset in range(node.lineno - 1, 0, -1):
line = source.split('\n')[offset - 1]
if line.strip():
# contains non-whitespace
try:
toks = _tokenize_line(line)
except (tokenize.TokenError, SyntaxError):
return None
# print(list(toks))
comment_toks = [x for x in list(toks) if x.type == tokenize.COMMENT]
if comment_toks:
for match in re.finditer(r"(#:\s*)(.*)", comment_toks[0].string):
if match.group(2):
return match.group(2).rstrip()
return None
return None
def _tokenize_line(line: str) -> List[tokenize.TokenInfo]:
"""
Tokenize a single line of Python source code.
:param line:
"""
def yielder() -> Iterator[str]:
yield line
return list(tokenize.generate_tokens(yielder().__next__))
[docs]class MultipleDocstringsWarning(UserWarning):
"""
Warning emitted when multiple docstrings are found for a single Enum member.
.. versionadded:: 0.8.0
:param member:
:param docstrings: The list of docstrings found for the member.
"""
#: The member with multiple docstrings.
member: Enum
#: The list of docstrings found for the member.
docstrings: Iterable[str]
def __init__(self, member: Enum, docstrings: Iterable[str] = ()):
self.member = member
self.docstrings = docstrings
[docs] def __str__(self) -> str:
member_full_name = '.'.join([
self.member.__class__.__module__,
self.member.__class__.__name__,
self.member.name,
])
return f"Found multiple docstrings for enum member <{member_full_name}>"
[docs]def document_enum(an_enum: EnumType) -> EnumType:
"""
Document all members of an enum by parsing a docstring from the Python source..
The docstring can be added in several ways:
#. A comment at the end the line, starting with ``doc:``:
.. code-block:: python
Running = 1 # doc: The system is running.
#. A comment on the previous line, starting with ``#:``. This is the format used by Sphinx.
.. code-block:: python
#: The system is running.
Running = 1
#. A string on the line *after* the attribute. This can be used for multiline docstrings.
.. code-block:: python
Running = 1
\"\"\"
The system is running.
Hello World
\"\"\"
If more than one docstring format is found for an enum member
a :exc:`MultipleDocstringsWarning` is emitted.
:param an_enum: An :class:`~enum.Enum` subclass
:type an_enum: :class:`enum.Enum`
:returns: The same object passed as ``an_enum``. This allows this function to be used as a decorator.
:rtype: :class:`enum.Enum`
.. versionchanged:: 0.8.0 Added support for other docstring formats and multiline docstrings.
"""
if not isinstance(an_enum, EnumMeta):
raise TypeError(f"'an_enum' must be an 'Enum', not {type(an_enum)}!")
if not INTERACTIVE:
return an_enum
func_source = dedent(inspect.getsource(an_enum))
func_source_tree = ast.parse(func_source)
assert len(func_source_tree.body) == 1
module_body = func_source_tree.body[0]
assert isinstance(module_body, ast.ClassDef)
class_body = module_body.body
for idx, node in enumerate(class_body):
targets = []
if isinstance(node, ast.Assign):
for t in node.targets:
assert isinstance(t, ast.Name)
targets.append(t.id)
elif isinstance(node, ast.AnnAssign):
assert isinstance(node.target, ast.Name)
targets.append(node.target.id)
else:
continue
assert isinstance(node, (ast.Assign, ast.AnnAssign))
# print(targets)
if idx + 1 == len(class_body):
next_node = None
else:
next_node = class_body[idx + 1]
docstring_candidates = []
if isinstance(next_node, ast.Expr):
# might be docstring
docstring_candidates.append(_docstring_from_expr(next_node))
# maybe no luck with """ docstring? look for EOL comment.
docstring_candidates.append(_docstring_from_eol_comment(func_source, node))
# check non-whitespace lines above for Sphinx-style comment.
docstring_candidates.append(_docstring_from_sphinx_comment(func_source, node))
docstring_candidates_nn = list(filter(None, docstring_candidates))
if len(docstring_candidates_nn) > 1:
# Multiple docstrings found, warn
warnings.warn(MultipleDocstringsWarning(getattr(an_enum, targets[0]), docstring_candidates_nn))
if docstring_candidates_nn:
docstring = docstring_candidates_nn[0]
for target in targets:
getattr(an_enum, target).__doc__ = docstring
return an_enum
[docs]def document_member(enum_member: Enum) -> None:
"""
Document a member of an enum by adding a comment to the end of the line that starts with ``doc:``.
:param enum_member: A member of an :class:`~enum.Enum` subclass
"""
if not isinstance(enum_member, Enum):
raise TypeError(f"'an_enum' must be an 'Enum', not {type(enum_member)}!")
if not INTERACTIVE:
return None
func_source = dedent(inspect.getsource(enum_member.__class__))
in_docstring = False
base_indent = None
for line in func_source.split('\n'):
indent, line = get_dedented_line(line)
if line.startswith("class") or not line:
continue
all_tokens = get_tokens(line)
base_indent = get_base_indent(base_indent, all_tokens, indent)
# print(all_tokens)
if enum_member.name not in line:
continue
if all_tokens[0][0] in pygments.token.Literal.String:
if all_tokens[0][1] in {'"""', "'''"}: # TODO: handle the other quotes appearing in docstring
in_docstring = not in_docstring
if all_tokens[0][0] in pygments.token.Name and in_docstring:
continue
elif all_tokens[0][0] not in pygments.token.Name:
continue
else:
if indent > base_indent: # type: ignore[operator]
continue
enum_vars, doc = parse_tokens(all_tokens)
for var in enum_vars:
# print(repr(var))
if not var.startswith('@'):
if var == enum_member.name:
enum_member.__doc__ = doc
return None
[docs]def parse_tokens(all_tokens: Iterable["pygments.Token"]) -> Tuple[List, Optional[str]]:
"""
Parse the tokens representing a line of code to identify Enum members and ``doc:`` comments.
:param all_tokens:
:return: A list of the Enum members' names, and the docstring for them.
"""
enum_vars = []
doc = None
comment = ''
for token in all_tokens:
if token[0] in pygments.token.Name:
enum_vars.append(token[1])
elif token[0] in pygments.token.Comment:
comment = token[1]
break
for match in re.finditer(r"(doc:\s*)([^#]*)(#|$)", comment):
if match.group(2):
doc = match.group(2).rstrip()
break
return enum_vars, doc
[docs]def get_base_indent(
base_indent: Optional[int],
all_tokens: Sequence[Sequence],
indent: int,
) -> Optional[int]:
"""
Determine the base level of indentation (i.e. one level of indentation in from the ``c`` of ``class``).
:param base_indent: The current base level of indentation
:param all_tokens:
:param indent: The current level of indentation
:returns: The base level of indentation
"""
if not base_indent:
if all_tokens[0][0] in pygments.token.Literal.String:
if all_tokens[0][1] in {'"""', "'''"}:
base_indent = indent
elif all_tokens[0][0] in pygments.token.Keyword:
base_indent = indent
elif all_tokens[0][0] in pygments.token.Name:
base_indent = indent
return base_indent
[docs]class DocumentedEnum(Enum):
"""
An enum where docstrings are automatically added to members from comments starting with ``doc:``.
.. note:: This class does not (yet) support the other docstring formats :deco:`~.document_enum` does.
"""
def __init__(self, value): # noqa: MAN001
document_member(self)
# super().__init__(value)
[docs]def get_dedented_line(line: str) -> Tuple[int, str]:
"""
Returns the line without indentation, and the amount of indentation.
:param line: A line of Python source code
"""
dedented_line = dedent(line)
indent = len(line) - len(dedented_line)
line = dedented_line.strip()
return indent, line