Fix #317 Help text short args

fire/helptext.py

@@ -33,6 +33,7 @@
 from __future__ import division
 from __future__ import print_function
 
+import collections
 import itertools
 import sys
 
@@ -172,9 +173,25 @@ def _DescriptionSection(component, info):
     return None
 
 
-def _CreateKeywordOnlyFlagItem(flag, docstring_info, spec):
+def _CreateKeywordOnlyFlagItem(flag, docstring_info, spec, short_arg):
   return _CreateFlagItem(
-      flag, docstring_info, spec, required=flag not in spec.kwonlydefaults)
+      flag, docstring_info, spec, required=flag not in spec.kwonlydefaults,
+      short_arg=short_arg)
+
+
+def _GetShortFlags(flags):
+  """Gets a list of single-character flags that uniquely identify a flag.
+
+  Args:
+    flags: list of strings representing flags
+
+  Returns:
+    List of single character short flags,
+    where the character occurred at the start of a flag once.
+  """
+  short_flags = [f[0] for f in flags]
+  short_flag_counts = collections.Counter(short_flags)
+  return [v for v in short_flags if short_flag_counts[v] == 1]
 
 
 def _ArgsAndFlagsSections(info, spec, metadata):
@@ -209,25 +226,47 @@ def _ArgsAndFlagsSections(info, spec, metadata):
           ('NOTES', 'You can also use flags syntax for POSITIONAL ARGUMENTS')
       )
 
+  unique_short_args = _GetShortFlags(args_with_defaults)
   positional_flag_items = [
-      _CreateFlagItem(flag, docstring_info, spec, required=False)
+      _CreateFlagItem(
+          flag, docstring_info, spec, required=False,
+          short_arg=flag[0] in unique_short_args
+      )
       for flag in args_with_defaults
   ]
+
+  unique_short_kwonly_flags = _GetShortFlags(spec.kwonlyargs)
   kwonly_flag_items = [
-      _CreateKeywordOnlyFlagItem(flag, docstring_info, spec)
+      _CreateKeywordOnlyFlagItem(
+          flag, docstring_info, spec,
+          short_arg=flag[0] in unique_short_kwonly_flags
+      )
       for flag in spec.kwonlyargs
   ]
   flag_items = positional_flag_items + kwonly_flag_items
 
   if spec.varkw:
     # Include kwargs documented via :key param:
-    flag_string = '--{name}'
     documented_kwargs = []
-    for flag in docstring_info.args or []:
+    flag_string = '--{name}'
+    short_flag_string = '-{short_name}, --{name}'
+
+    # add short flags if possible
+    flags = docstring_info.args or []
+    flag_names = [f.name for f in flags]
+    unique_short_flags = _GetShortFlags(flag_names)
+    for flag in flags:
       if isinstance(flag, docstrings.KwargInfo):
+        if flag.name[0] in unique_short_flags:
+          flag_string = short_flag_string.format(
+              name=flag.name, short_name=flag.name[0]
+          )
+        else:
+          flag_string = flag_string.format(name=flag.name)
+
         flag_item = _CreateFlagItem(
             flag.name, docstring_info, spec,
-            flag_string=flag_string.format(name=flag.name))
+            flag_string=flag_string)
         documented_kwargs.append(flag_item)
     if documented_kwargs:
       # Separate documented kwargs from other flags using a message
@@ -422,7 +461,7 @@ def _CreateArgItem(arg, docstring_info, spec):
 
 
 def _CreateFlagItem(flag, docstring_info, spec, required=False,
-                    flag_string=None):
+                    flag_string=None, short_arg=False):
   """Returns a string describing a flag using docstring and FullArgSpec info.
 
   Args:
@@ -434,6 +473,7 @@ def _CreateFlagItem(flag, docstring_info, spec, required=False,
     required: Whether the flag is required.
     flag_string: If provided, use this string for the flag, rather than
       constructing one from the flag name.
+    short_arg: Whether the flag has a short variation or not.
   Returns:
     A string to be used in constructing the help screen for the function.
   """
@@ -455,6 +495,8 @@ def _CreateFlagItem(flag, docstring_info, spec, required=False,
         flag_name_upper=formatting.Underline(flag.upper()))
   if required:
     flag_string += ' (required)'
+  if short_arg:
+    flag_string = '-{short_flag}, '.format(short_flag=flag[0]) + flag_string
 
   arg_type = _GetArgType(flag, spec)
   arg_default = _GetArgDefault(flag, spec)

fire/helptext_test.py

@@ -81,7 +81,9 @@ def testHelpTextFunctionWithDefaults(self):
     self.assertIn('NAME\n    triple', help_screen)
     self.assertIn('SYNOPSIS\n    triple <flags>', help_screen)
     self.assertNotIn('DESCRIPTION', help_screen)
-    self.assertIn('FLAGS\n    --count=COUNT\n        Default: 0', help_screen)
+    self.assertIn(
+        'FLAGS\n    -c, --count=COUNT\n        Default: 0',
+        help_screen)
     self.assertNotIn('NOTES', help_screen)
 
   def testHelpTextFunctionWithLongDefaults(self):
@@ -93,7 +95,7 @@ def testHelpTextFunctionWithLongDefaults(self):
     self.assertIn('SYNOPSIS\n    text <flags>', help_screen)
     self.assertNotIn('DESCRIPTION', help_screen)
     self.assertIn(
-        'FLAGS\n    --string=STRING\n'
+        'FLAGS\n    -s, --string=STRING\n'
         '        Default: \'0001020304050607080910'
         '1112131415161718192021222324252627282...',
         help_screen)
@@ -121,7 +123,7 @@ def testHelpTextFunctionWithKwargsAndDefaults(self):
     self.assertIn('SYNOPSIS\n    text ARG1 ARG2 <flags>', help_screen)
     self.assertIn('DESCRIPTION\n    Function with kwarg', help_screen)
     self.assertIn(
-        'FLAGS\n    --opt=OPT\n        Default: True\n'
+        'FLAGS\n    -o, --opt=OPT\n        Default: True\n'
         '    The following flags are also accepted.'
         '\n    --arg3\n        Description of arg3.\n    '
         'Additional undocumented flags may also be accepted.',
@@ -140,7 +142,7 @@ def testHelpTextFunctionWithDefaultsAndTypes(self):
     self.assertIn('SYNOPSIS\n    double <flags>', help_screen)
     self.assertIn('DESCRIPTION', help_screen)
     self.assertIn(
-        'FLAGS\n    --count=COUNT\n        Type: float\n        Default: 0',
+        'FLAGS\n    -c, --count=COUNT\n        Type: float\n        Default: 0',
         help_screen)
     self.assertNotIn('NOTES', help_screen)
 
@@ -157,7 +159,7 @@ def testHelpTextFunctionWithTypesAndDefaultNone(self):
     self.assertIn('SYNOPSIS\n    get_int <flags>', help_screen)
     self.assertNotIn('DESCRIPTION', help_screen)
     self.assertIn(
-        'FLAGS\n    --value=VALUE\n'
+        'FLAGS\n    -v, --value=VALUE\n'
         '        Type: Optional[int]\n        Default: None',
         help_screen)
     self.assertNotIn('NOTES', help_screen)
@@ -285,7 +287,7 @@ def testHelpTextKeywordOnlyArgumentsWithDefault(self):
     output = helptext.HelpText(
         component=component, trace=trace.FireTrace(component, 'with_default'))
     self.assertIn('NAME\n    with_default', output)
-    self.assertIn('FLAGS\n    --x=X', output)
+    self.assertIn('FLAGS\n    -x, --x=X', output)
 
   @testutils.skipIf(
       six.PY2, 'Python 2 does not support keyword-only arguments.')
@@ -294,7 +296,7 @@ def testHelpTextKeywordOnlyArgumentsWithoutDefault(self):
     output = helptext.HelpText(
         component=component, trace=trace.FireTrace(component, 'double'))
     self.assertIn('NAME\n    double', output)
-    self.assertIn('FLAGS\n    --count=COUNT (required)', output)
+    self.assertIn('FLAGS\n    -c, --count=COUNT (required)', output)
 
   def testHelpScreen(self):
     component = tc.ClassWithDocstring()
@@ -362,7 +364,7 @@ def testHelpScreenForFunctionFunctionWithDefaultArgs(self):
         Returns the input multiplied by 2.
 
     FLAGS
-        --count=COUNT
+        -c, --count=COUNT
             Default: 0
             Input number that you want to double."""
     self.assertEqual(textwrap.dedent(expected_output).strip(),
@@ -377,7 +379,8 @@ def testHelpTextUnderlineFlag(self):
         formatting.Bold('SYNOPSIS') + '\n    triple <flags>',
         help_screen)
     self.assertIn(
-        formatting.Bold('FLAGS') + '\n    --' + formatting.Underline('count'),
+        formatting.Bold('FLAGS') + '\n    -c, --' +
+        formatting.Underline('count'),
         help_screen)
 
   def testHelpTextBoldCommandName(self):
@@ -423,6 +426,21 @@ def testHelpTextNameSectionCommandWithSeparatorVerbose(self):
     self.assertIn('double -', help_screen)
     self.assertIn('double - -', help_screen)
 
+  def testHelpTextMultipleKeywoardArgumentsWithShortArgs(self):
+    component = tc.fn_with_multiple_defaults
+    t = trace.FireTrace(component, name='shortargs')
+    help_screen = helptext.HelpText(component, t)
+    self.assertIn(formatting.Bold('NAME') + '\n    shortargs', help_screen)
+    self.assertIn(
+        formatting.Bold('SYNOPSIS') + '\n    shortargs <flags>',
+        help_screen)
+    self.assertIn(
+        formatting.Bold('FLAGS') + '\n    -f, --first',
+        help_screen)
+    self.assertIn('\n    --last', help_screen)
+    self.assertIn('\n    --late', help_screen)
+
+
 
 class UsageTest(testutils.BaseTestCase):
 

fire/test_components.py

@@ -563,3 +563,13 @@ def fn_with_kwarg_and_defaults(arg1, arg2, opt=True, **kwargs):
   del arg1, arg2, opt
   return kwargs.get('arg3')
 # pylint: enable=g-doc-args,g-doc-return-or-yield
+
+def fn_with_multiple_defaults(first='first', last='last', late='late'):
+  """Function with kwarg and defaults.
+
+  :key first: Description of first.
+  :key last: Description of last.
+  :key late: Description of late.
+  """
+  del last, late
+  return first