6
6
# LICENSE file in the root directory of this source tree.
7
7
8
8
import abc
9
+ import argparse
9
10
import ast
11
+ import inspect
10
12
from dataclasses import dataclass
11
- from typing import Dict , List , Optional , Tuple , cast
13
+ from typing import Dict , List , Optional , Tuple , cast , Callable
12
14
13
15
from docstring_parser import parse
14
16
from pyre_extensions import none_throws
18
20
# pyre-ignore-all-errors[16]
19
21
20
22
21
- def get_arg_names (app_specs_func_def : ast .FunctionDef ) -> List [str ]:
22
- arg_names = []
23
- fn_args = app_specs_func_def .args
24
- for arg_def in fn_args .args :
25
- arg_names .append (arg_def .arg )
26
- if fn_args .vararg :
27
- arg_names .append (fn_args .vararg .arg )
28
- for arg in fn_args .kwonlyargs :
29
- arg_names .append (arg .arg )
30
- return arg_names
23
+ def _get_default_arguments_descriptions (fn : Callable [..., object ]) -> Dict [str , str ]:
24
+ parameters = inspect .signature (fn ).parameters
25
+ args_decs = {}
26
+ for parameter_name in parameters .keys ():
27
+ # The None or Empty string values getting ignored during help command by argparse
28
+ args_decs [parameter_name ] = " "
29
+ return args_decs
31
30
32
31
33
- def parse_fn_docstring (func_description : str ) -> Tuple [str , Dict [str , str ]]:
32
+ class TorchXArgumentHelpFormatter (argparse .HelpFormatter ):
33
+ """Help message formatter which adds default values and required to argument help.
34
+
35
+ If the argument is required, the class appends `(required)` at the end of the help message.
36
+ If the argument has default value, the class appends `(default: $DEFAULT)` at the end.
37
+ The formatter is designed to be used only for the torchx components functions.
38
+ These functions do not have both required and default arguments.
34
39
"""
35
- Given a docstring in a google-style format, returns the function description and
36
- description of all arguments.
37
- See: https://sphinxcontrib-napoleon.readthedocs.io/en/latest/example_google.html
40
+
41
+ def _get_help_string (self , action : argparse .Action ) -> str :
42
+ help = action .help or ""
43
+ # Only `--help` will have be SUPPRESS, so we ignore it
44
+ if action .default is argparse .SUPPRESS :
45
+ return help
46
+ if action .required :
47
+ help += " (required)"
48
+ else :
49
+ help += f" (default: { action .default } )"
50
+ return help
51
+
52
+
53
+ def get_fn_docstring (fn : Callable [..., object ]) -> Tuple [str , Dict [str , str ]]:
38
54
"""
39
- args_description = {}
55
+ Parses the function and arguments description from the provided function. Docstring should be in
56
+ `google-style format <https://sphinxcontrib-napoleon.readthedocs.io/en/latest/example_google.html>`_
57
+
58
+ If function has no docstring, the function description will be the name of the function, TIP
59
+ on how to improve the help message and arguments descriptions will be names of the arguments.
60
+
61
+ The arguments that are not present in the docstring will contain default/required information
62
+
63
+ Args:
64
+ fn: Function with or without docstring
65
+
66
+ Returns:
67
+ function description, arguments description where key is the name of the argument and value
68
+ if the description
69
+ """
70
+ default_fn_desc = f"""{ fn .__name__ } TIP: improve this help string by adding a docstring
71
+ to your component (see: https://pytorch.org/torchx/latest/component_best_practices.html)"""
72
+ args_description = _get_default_arguments_descriptions (fn )
73
+ func_description = inspect .getdoc (fn )
74
+ if not func_description :
75
+ return default_fn_desc , args_description
40
76
docstring = parse (func_description )
41
77
for param in docstring .params :
42
78
args_description [param .arg_name ] = param .description
43
- short_func_description = docstring .short_description
44
- return (short_func_description or "" , args_description )
45
-
46
-
47
- def _get_fn_docstring (
48
- source : str , function_name : str
49
- ) -> Optional [Tuple [str , Dict [str , str ]]]:
50
- module = ast .parse (source )
51
- for expr in module .body :
52
- if type (expr ) == ast .FunctionDef :
53
- func_def = cast (ast .FunctionDef , expr )
54
- if func_def .name == function_name :
55
- docstring = ast .get_docstring (func_def )
56
- if not docstring :
57
- return None
58
- return parse_fn_docstring (docstring )
59
- return None
60
-
61
-
62
- def get_short_fn_description (path : str , function_name : str ) -> Optional [str ]:
63
- source = read_conf_file (path )
64
- docstring = _get_fn_docstring (source , function_name )
65
- if not docstring :
66
- return None
67
- return docstring [0 ]
79
+ short_func_description = docstring .short_description or default_fn_desc
80
+ if docstring .long_description :
81
+ short_func_description += " ..."
82
+ return (short_func_description or default_fn_desc , args_description )
68
83
69
84
70
85
@dataclass
@@ -91,38 +106,6 @@ def _gen_linter_message(self, description: str, lineno: int) -> LinterMessage:
91
106
)
92
107
93
108
94
- class TorchxDocstringValidator (TorchxFunctionValidator ):
95
- def validate (self , app_specs_func_def : ast .FunctionDef ) -> List [LinterMessage ]:
96
- """
97
- Validates the docstring of the `get_app_spec` function. Criteria:
98
- * There mast be google-style docstring
99
- * If there are more than zero arguments, there mast be a `Args:` section defined
100
- with all arguments included.
101
- """
102
- docsting = ast .get_docstring (app_specs_func_def )
103
- lineno = app_specs_func_def .lineno
104
- if not docsting :
105
- desc = (
106
- f"`{ app_specs_func_def .name } ` is missing a Google Style docstring, please add one. "
107
- "For more information on the docstring format see: "
108
- "https://sphinxcontrib-napoleon.readthedocs.io/en/latest/example_google.html"
109
- )
110
- return [self ._gen_linter_message (desc , lineno )]
111
-
112
- arg_names = get_arg_names (app_specs_func_def )
113
- _ , docstring_arg_defs = parse_fn_docstring (docsting )
114
- missing_args = [
115
- arg_name for arg_name in arg_names if arg_name not in docstring_arg_defs
116
- ]
117
- if len (missing_args ) > 0 :
118
- desc = (
119
- f"`{ app_specs_func_def .name } ` not all function arguments are present"
120
- f" in the docstring. Missing args: { missing_args } "
121
- )
122
- return [self ._gen_linter_message (desc , lineno )]
123
- return []
124
-
125
-
126
109
class TorchxFunctionArgsValidator (TorchxFunctionValidator ):
127
110
def validate (self , app_specs_func_def : ast .FunctionDef ) -> List [LinterMessage ]:
128
111
linter_errors = []
@@ -149,7 +132,6 @@ def _validate_arg_def(
149
132
)
150
133
]
151
134
if isinstance (arg_def .annotation , ast .Name ):
152
- # TODO(aivanou): add support for primitive type check
153
135
return []
154
136
complex_type_def = cast (ast .Subscript , none_throws (arg_def .annotation ))
155
137
if complex_type_def .value .id == "Optional" :
@@ -239,12 +221,6 @@ class TorchFunctionVisitor(ast.NodeVisitor):
239
221
Visitor that finds the component_function and runs registered validators on it.
240
222
Current registered validators:
241
223
242
- * TorchxDocstringValidator - validates the docstring of the function.
243
- Criteria:
244
- * There format should be google-python
245
- * If there are more than zero arguments defined, there
246
- should be obligatory `Args:` section that describes each argument on a new line.
247
-
248
224
* TorchxFunctionArgsValidator - validates arguments of the function.
249
225
Criteria:
250
226
* Each argument should be annotated with the type
@@ -260,7 +236,6 @@ class TorchFunctionVisitor(ast.NodeVisitor):
260
236
261
237
def __init__ (self , component_function_name : str ) -> None :
262
238
self .validators = [
263
- TorchxDocstringValidator (),
264
239
TorchxFunctionArgsValidator (),
265
240
TorchxReturnValidator (),
266
241
]
0 commit comments