Make docstring optional (#259)

Summary:
Pull Request resolved: #259

* Refactor docstring functions: combines two functions that retrieve docstring into one
* Make docstring optional
* Remove docstring validator

Differential Revision: D31671125

fbshipit-source-id: 2fb71f6b98e212700479003ca4d15a01cec0e571

Author: aivanou

torchx/specs/api.py

@@ -28,8 +28,7 @@
 )
 
 import yaml
-from pyre_extensions import none_throws
-from torchx.specs.file_linter import parse_fn_docstring
+from torchx.specs.file_linter import get_fn_docstring
 from torchx.util.types import decode_from_string, decode_optional, is_bool, is_primitive
 
 
@@ -748,22 +747,21 @@ def get_argparse_param_type(parameter: inspect.Parameter) -> Callable[[str], obj
         return str
 
 
-def _create_args_parser(
-    fn_name: str,
-    parameters: Mapping[str, inspect.Parameter],
-    function_desc: str,
-    args_desc: Dict[str, str],
-) -> argparse.ArgumentParser:
+def _create_args_parser(app_fn: Callable[..., AppDef]) -> argparse.ArgumentParser:
+    parameters = inspect.signature(app_fn).parameters
+    function_desc, args_desc = get_fn_docstring(app_fn)
     script_parser = argparse.ArgumentParser(
-        prog=f"torchx run ...torchx_params... {fn_name} ",
+        prog=f"torchx run <<torchx_params>> {app_fn.__name__} ",
         description=f"App spec: {function_desc}",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
     )
 
     remainder_arg = []
 
     for param_name, parameter in parameters.items():
+        param_desc = args_desc[parameter.name]
         args: Dict[str, Any] = {
-            "help": args_desc[param_name],
+            "help": param_desc,
             "type": get_argparse_param_type(parameter),
         }
         if parameter.default != inspect.Parameter.empty:
@@ -788,20 +786,15 @@ def _create_args_parser(
 def _get_function_args(
     app_fn: Callable[..., AppDef], app_args: List[str]
 ) -> Tuple[List[object], List[str], Dict[str, object]]:
-    docstring = none_throws(inspect.getdoc(app_fn))
-    function_desc, args_desc = parse_fn_docstring(docstring)
-
-    parameters = inspect.signature(app_fn).parameters
-    script_parser = _create_args_parser(
-        app_fn.__name__, parameters, function_desc, args_desc
-    )
+    script_parser = _create_args_parser(app_fn)
 
     parsed_args = script_parser.parse_args(app_args)
 
     function_args = []
     var_arg = []
     kwargs = {}
 
+    parameters = inspect.signature(app_fn).parameters
     for param_name, parameter in parameters.items():
         arg_value = getattr(parsed_args, param_name)
         parameter_type = parameter.annotation

torchx/specs/file_linter.py

@@ -7,8 +7,9 @@
 
 import abc
 import ast
+import inspect
 from dataclasses import dataclass
-from typing import Dict, List, Optional, Tuple, cast
+from typing import Dict, List, Optional, Tuple, cast, Callable
 
 from docstring_parser import parse
 from pyre_extensions import none_throws
@@ -30,41 +31,29 @@ def get_arg_names(app_specs_func_def: ast.FunctionDef) -> List[str]:
     return arg_names
 
 
-def parse_fn_docstring(func_description: str) -> Tuple[str, Dict[str, str]]:
+def _get_default_arguments_descriptions(fn: Callable[..., object]) -> Dict[str, str]:
+    parameters = inspect.signature(fn).parameters
+    args_decs = {}
+    for parameter_name in parameters.keys():
+        args_decs[parameter_name] = parameter_name
+    return args_decs
+
+
+def get_fn_docstring(fn: Callable[..., object]) -> Tuple[str, Dict[str, str]]:
     """
     Given a docstring in a google-style format, returns the function description and
     description of all arguments.
     See: https://sphinxcontrib-napoleon.readthedocs.io/en/latest/example_google.html
     """
-    args_description = {}
+    args_description = _get_default_arguments_descriptions(fn)
+    func_description = inspect.getdoc(fn)
+    if not func_description:
+        return fn.__name__, args_description
     docstring = parse(func_description)
     for param in docstring.params:
         args_description[param.arg_name] = param.description
     short_func_description = docstring.short_description
-    return (short_func_description or "", args_description)
-
-
-def _get_fn_docstring(
-    source: str, function_name: str
-) -> Optional[Tuple[str, Dict[str, str]]]:
-    module = ast.parse(source)
-    for expr in module.body:
-        if type(expr) == ast.FunctionDef:
-            func_def = cast(ast.FunctionDef, expr)
-            if func_def.name == function_name:
-                docstring = ast.get_docstring(func_def)
-                if not docstring:
-                    return None
-                return parse_fn_docstring(docstring)
-    return None
-
-
-def get_short_fn_description(path: str, function_name: str) -> Optional[str]:
-    source = read_conf_file(path)
-    docstring = _get_fn_docstring(source, function_name)
-    if not docstring:
-        return None
-    return docstring[0]
+    return (short_func_description or fn.__name__, args_description)
 
 
 @dataclass
@@ -91,38 +80,6 @@ def _gen_linter_message(self, description: str, lineno: int) -> LinterMessage:
         )
 
 
-class TorchxDocstringValidator(TorchxFunctionValidator):
-    def validate(self, app_specs_func_def: ast.FunctionDef) -> List[LinterMessage]:
-        """
-        Validates the docstring of the `get_app_spec` function. Criteria:
-        * There mast be google-style docstring
-        * If there are more than zero arguments, there mast be a `Args:` section defined
-            with all arguments included.
-        """
-        docsting = ast.get_docstring(app_specs_func_def)
-        lineno = app_specs_func_def.lineno
-        if not docsting:
-            desc = (
-                f"`{app_specs_func_def.name}` is missing a Google Style docstring, please add one. "
-                "For more information on the docstring format see: "
-                "https://sphinxcontrib-napoleon.readthedocs.io/en/latest/example_google.html"
-            )
-            return [self._gen_linter_message(desc, lineno)]
-
-        arg_names = get_arg_names(app_specs_func_def)
-        _, docstring_arg_defs = parse_fn_docstring(docsting)
-        missing_args = [
-            arg_name for arg_name in arg_names if arg_name not in docstring_arg_defs
-        ]
-        if len(missing_args) > 0:
-            desc = (
-                f"`{app_specs_func_def.name}` not all function arguments are present"
-                f" in the docstring. Missing args: {missing_args}"
-            )
-            return [self._gen_linter_message(desc, lineno)]
-        return []
-
-
 class TorchxFunctionArgsValidator(TorchxFunctionValidator):
     def validate(self, app_specs_func_def: ast.FunctionDef) -> List[LinterMessage]:
         linter_errors = []
@@ -149,7 +106,6 @@ def _validate_arg_def(
                 )
             ]
         if isinstance(arg_def.annotation, ast.Name):
-            # TODO(aivanou): add support for primitive type check
             return []
         complex_type_def = cast(ast.Subscript, none_throws(arg_def.annotation))
         if complex_type_def.value.id == "Optional":
@@ -239,12 +195,6 @@ class TorchFunctionVisitor(ast.NodeVisitor):
     Visitor that finds the component_function and runs registered validators on it.
     Current registered validators:
 
-    * TorchxDocstringValidator - validates the docstring of the function.
-        Criteria:
-          * There format should be google-python
-          * If there are more than zero arguments defined, there
-            should be obligatory `Args:` section that describes each argument on a new line.
-
     * TorchxFunctionArgsValidator - validates arguments of the function.
         Criteria:
           * Each argument should be annotated with the type
@@ -260,7 +210,6 @@ class TorchFunctionVisitor(ast.NodeVisitor):
 
     def __init__(self, component_function_name: str) -> None:
         self.validators = [
-            TorchxDocstringValidator(),
             TorchxFunctionArgsValidator(),
             TorchxReturnValidator(),
         ]

torchx/specs/finder.py

@@ -17,7 +17,7 @@
 
 from pyre_extensions import none_throws
 from torchx.specs import AppDef
-from torchx.specs.file_linter import get_short_fn_description, validate
+from torchx.specs.file_linter import get_fn_docstring, validate
 from torchx.util import entrypoints
 from torchx.util.io import read_conf_file
 
@@ -47,7 +47,7 @@ class _Component:
     """
 
     name: str
-    description: Optional[str]
+    description: str
     fn_name: str
     fn: Callable[..., AppDef]
     validation_errors: List[str]
@@ -119,9 +119,10 @@ def _get_components_from_dir(
         search_pattern = os.path.join(search_dir, "**", "*.py")
         component_defs = []
         for filepath in glob.glob(search_pattern, recursive=True):
-            module = self._try_load_module(
-                self._get_module_name(filepath, search_dir, base_module)
-            )
+            module_name = self._get_module_name(filepath, search_dir, base_module)
+            if module_name.startswith("torchx.components.base"):
+                continue
+            module = self._try_load_module(module_name)
             defs = self._get_components_from_module(base_module, module)
             component_defs += defs
         return component_defs
@@ -146,7 +147,7 @@ def _get_components_from_module(
         module_path = os.path.abspath(module.__file__)
         for function_name, function in functions:
             linter_errors = validate(module_path, function_name)
-            component_desc = get_short_fn_description(module_path, function_name)
+            component_desc, _ = get_fn_docstring(function)
             component_def = _Component(
                 name=self._get_component_name(
                     base_module, module.__name__, function_name
@@ -193,7 +194,6 @@ def find(self) -> List[_Component]:
         validation_errors = self._get_validation_errors(
             self._filepath, self._function_name
         )
-        fn_desc = get_short_fn_description(self._filepath, self._function_name)
 
         file_source = read_conf_file(self._filepath)
         namespace = globals()
@@ -203,6 +203,7 @@ def find(self) -> List[_Component]:
                 f"Function {self._function_name} does not exist in file {self._filepath}"
             )
         app_fn = namespace[self._function_name]
+        fn_desc, _ = get_fn_docstring(app_fn)
         return [
             _Component(
                 name=f"{self._filepath}:{self._function_name}",

torchx/specs/test/api_test.py

@@ -5,6 +5,7 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+import argparse
 import sys
 import unittest
 from dataclasses import asdict
@@ -33,6 +34,7 @@
     make_app_handle,
     parse_app_handle,
     runopts,
+    _create_args_parser,
 )
 
 
@@ -463,11 +465,6 @@ def _test_complex_fn(
         app_name: AppDef name
         containers: List of containers
         roles_scripts: Dict role_name -> role_script
-        num_cpus: List of cpus per role
-        num_gpus: Dict role_name -> gpus used for role
-        nnodes: Num replicas per role
-        first_arg: First argument to the user script
-        roles_args: Roles args
     """
     num_roles = len(roles_scripts)
     if not num_cpus:
@@ -710,3 +707,28 @@ def test_varargs_only_arg_first(self) -> None:
             _TEST_VAR_ARGS_FIRST,
             (("fooval", "--foo", "barval", "arg1", "arg2"), "asdf"),
         )
+
+    def _get_argument_help(
+        self, parser: argparse.ArgumentParser, name: str
+    ) -> Optional[str]:
+        actions = parser._actions
+        for action in actions:
+            if action.dest == name:
+                return action.help
+        return None
+
+    def test_argparster_complex_fn_partial(self) -> None:
+        parser = _create_args_parser(_test_complex_fn)
+        self.assertEqual("AppDef name", self._get_argument_help(parser, "app_name"))
+        self.assertEqual(
+            "List of containers", self._get_argument_help(parser, "containers")
+        )
+        self.assertEqual(
+            "Dict role_name -> role_script",
+            self._get_argument_help(parser, "roles_scripts"),
+        )
+        self.assertEqual("num_cpus", self._get_argument_help(parser, "num_cpus"))
+        self.assertEqual("num_gpus", self._get_argument_help(parser, "num_gpus"))
+        self.assertEqual("nnodes", self._get_argument_help(parser, "nnodes"))
+        self.assertEqual("first_arg", self._get_argument_help(parser, "first_arg"))
+        self.assertEqual("roles_args", self._get_argument_help(parser, "roles_args"))