Skip to content
Navigation Menu
{{ message }}
-
Notifications
You must be signed in to change notification settings - Fork 8
Expand file tree
/
Copy pathdocumentation.py
More file actions
422 lines (308 loc) · 10.8 KB
/
Copy pathdocumentation.py
File metadata and controls
422 lines (308 loc) · 10.8 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
#!/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-untyped]
from pygments.lexers.python import PythonLexer # type: ignore[import-untyped]
__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)
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__))
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
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}>"
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
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
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
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
class DocumentedEnum(Enum): # noqa: PRM002
"""
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)
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
You can’t perform that action at this time.