Generating flags information for function objects

PiperOrigin-RevId: 239247774
Change-Id: I2dc83a9006f94239cfc5b267c434304e4317f39a

fire/helputils.py

@@ -106,6 +106,20 @@ def _GetFields(trace=None):
   ]
 
 
+def GetArgsAngFlags(component):
+  """Returns all types of arguments and flags of a component."""
+  spec = inspectutils.GetFullArgSpec(component)
+  args = spec.args
+  if spec.defaults is None:
+    num_defaults = 0
+  else:
+    num_defaults = len(spec.defaults)
+  args_with_no_defaults = args[:len(args) - num_defaults]
+  args_with_defaults = args[len(args) - num_defaults:]
+  flags = args_with_defaults + spec.kwonlyargs
+  return args_with_no_defaults, args_with_defaults, flags
+
+
 def HelpString(component, trace=None, verbose=False):
   """Returns a help string for a supplied component.
 
@@ -220,15 +234,8 @@ def HelpTextForFunction(component, info, trace=None, verbose=False):
   spec = inspectutils.GetFullArgSpec(component)
   args = spec.args
 
-  if spec.defaults is None:
-    num_defaults = 0
-  else:
-    num_defaults = len(spec.defaults)
-  args_with_no_defaults = args[:len(args) - num_defaults]
-
-  # TODO(joejoevictor): Generate flag section using these
-  # args_with_defaults = args[len(args) - num_defaults:]
-  # flags = args_with_defaults + spec.kwonlyargs
+  args_with_no_defaults, args_with_defaults, flags = GetArgsAngFlags(component)
+  del args_with_defaults
 
   output_template = """NAME
     {name_section}
@@ -249,8 +256,18 @@ def HelpTextForFunction(component, info, trace=None, verbose=False):
   name_section = name_section_template.format(
       current_command=current_command, command_summary=command_summary_str)
 
-  items = [arg.upper() for arg in args_with_no_defaults]
-  args_and_flags = ' '.join(items)
+  args_and_flags = ''
+  if args_with_no_defaults:
+    items = [arg.upper() for arg in args_with_no_defaults]
+    args_and_flags = ' '.join(items)
+
+  synopsis_flag_template = '[--{flag_name}={flag_name_upper}]'
+  if flags:
+    items = [
+        synopsis_flag_template.format(
+            flag_name=flag, flag_name_upper=flag.upper()) for flag in flags
+    ]
+    args_and_flags = args_and_flags + ' '.join(items)
 
   # Synopsis section
   synopsis_section_template = '{current_command} {args_and_flags}'
@@ -266,7 +283,6 @@ def HelpTextForFunction(component, info, trace=None, verbose=False):
   args_and_flags_section = ''
 
   # Positional arguments and flags section
-
   pos_arg_template = """
 POSITIONAL ARGUMENTS
 {items}
@@ -286,6 +302,25 @@ def HelpTextForFunction(component, info, trace=None, verbose=False):
     args_and_flags_section += pos_arg_template.format(
         items='\n'.join(pos_arg_items).rstrip('\n'))
 
+  flags_template = """
+FLAGS
+{items}
+"""
+  flag_items = []
+  for flag in flags:
+    item_template = '    --{flag_name}\n        {flag_description}\n'
+    flag_description = None
+    for arg_in_docstring in info['docstring_info'].args:
+      if arg_in_docstring.name == flag:
+        flag_description = arg_in_docstring.description
+
+    item = item_template.format(
+        flag_name=flag, flag_description=flag_description)
+    flag_items.append(item)
+  if flag_items:
+    args_and_flags_section += flags_template.format(
+        items='\n'.join(flag_items).rstrip('\n'))
+
   return output_template.format(
       name_section=name_section,
       synopsis_section=synopsis_section,

fire/helputils_test.py

@@ -164,7 +164,7 @@ def testHelpScreen(self):
 """
     self.assertEqual(textwrap.dedent(expected_output).lstrip('\n'), help_output)
 
-  def testHelpScreenWithLineBreak(self):
+  def testHelpScreenForFunction_docstringWithLineBreak(self):
     component = tc.ClassWithMultilineDocstring.example_generator
     t = trace.FireTrace(component, name='example_generator')
     info = inspectutils.Info(component)
@@ -189,6 +189,31 @@ def testHelpScreenWithLineBreak(self):
     """
     self.assertEqual(textwrap.dedent(expected_output).lstrip('\n'), help_output)
 
+  def testHelpScreenForFunction_functionWithDefaultArgs(self):
+    component = tc.WithDefaults().double
+    t = trace.FireTrace(component, name='double')
+    info = inspectutils.Info(component)
+    info['docstring_info'] = docstrings.parse(info['docstring'])
+    help_output = helputils.HelpText(component, info, t)
+    expected_output = """
+    NAME
+        double - Returns the input multiplied by 2.
+
+    SYNOPSIS
+        double [--count=COUNT]
+
+    DESCRIPTION
+        Returns the input multiplied by 2.
+
+    FLAGS
+        --count
+            Input number that you want to double.
+
+    NOTES
+        You could also use flags syntax for POSITIONAL ARGUMENTS
+    """
+    self.assertEqual(textwrap.dedent(expected_output).lstrip('\n'), help_output)
+
 
 class UsageTest(testutils.BaseTestCase):
 

fire/test_components.py

@@ -87,8 +87,17 @@ def triple(self, count):
 
 
 class WithDefaults(object):
+  """Class with functions that have default arguments."""
 
   def double(self, count=0):
+    """Returns the input multiplied by 2.
+
+    Args:
+      count: Input number that you want to double.
+
+    Returns:
+      A number that is the double of count.s
+    """
     return 2 * count
 
   def triple(self, count=0):