bug fix for help from typing package, add test for optional and union

fire/helptext.py

@@ -31,6 +31,8 @@
 
 import collections
 import itertools
+import sys
+import typing
 
 from fire import completion
 from fire import custom_descriptions
@@ -488,7 +490,7 @@ def _CreateFlagItem(flag, docstring_info, spec, required=False,
 
   # We need to handle the case where there is a default of None, but otherwise
   # the argument has another type.
-  if arg_default == 'None':
+  if arg_default == 'None' and not arg_type.startswith('Optional'):
     arg_type = f'Optional[{arg_type}]'
 
   arg_type = f'Type: {arg_type}' if arg_type else ''
@@ -522,11 +524,11 @@ def _GetArgType(arg, spec):
   if arg in spec.annotations:
     arg_type = spec.annotations[arg]
     try:
+      if isinstance(arg_type, typing._GenericAlias):
+        arg_type = repr(arg_type).replace('typing.', '')
+        return arg_type
       return arg_type.__qualname__
     except AttributeError:
-      # Some typing objects, such as typing.Union do not have either a __name__
-      # or __qualname__ attribute.
-      # repr(typing.Union[int, str]) will return ': typing.Union[int, str]'
       return repr(arg_type)
   return ''
 

fire/helptext_test.py

@@ -152,6 +152,36 @@ def testHelpTextFunctionWithTypesAndDefaultNone(self):
         help_screen)
     self.assertNotIn('NOTES', help_screen)
 
+  def testHelpTextFunctionWithTypesAndDefaultNoneFromTypingOptional(self):
+    component = (
+        tc.py3.WithDefaultsAndTypes().typing_optional_get_int)  # pytype: disable=module-attr
+    help_screen = helptext.HelpText(
+        component=component,
+        trace=trace.FireTrace(component, name='get_int'))
+    self.assertIn('NAME\n    get_int', help_screen)
+    self.assertIn('SYNOPSIS\n    get_int <flags>', help_screen)
+    self.assertNotIn('DESCRIPTION', help_screen)
+    self.assertIn(
+        'FLAGS\n    -v, --value=VALUE\n'
+        '        Type: Optional[int]\n        Default: None',
+        help_screen)
+    self.assertNotIn('NOTES', help_screen)
+
+  def testHelpTextFunctionWithTypesAndDefaultNoneFromTypingUnion(self):
+    component = (
+        tc.py3.WithDefaultsAndTypes().typing_union_get_int)  # pytype: disable=module-attr
+    help_screen = helptext.HelpText(
+        component=component,
+        trace=trace.FireTrace(component, name='get_int'))
+    self.assertIn('NAME\n    get_int', help_screen)
+    self.assertIn('SYNOPSIS\n    get_int <flags>', help_screen)
+    self.assertNotIn('DESCRIPTION', help_screen)
+    self.assertIn(
+        'FLAGS\n    -v, --value=VALUE\n'
+        '        Type: Optional[Union[int, str]]\n        Default: None',
+        help_screen)
+    self.assertNotIn('NOTES', help_screen)
+
   def testHelpTextFunctionWithTypes(self):
     component = tc.py3.WithTypes().double  # pytype: disable=module-attr
     help_screen = helptext.HelpText(

fire/test_components_py3.py

@@ -16,7 +16,7 @@
 
 import asyncio
 import functools
-from typing import Tuple
+from typing import Tuple, Optional, Union
 
 
 # pylint: disable=keyword-arg-before-vararg
@@ -99,3 +99,9 @@ def double(self, count: float = 0) -> float:
 
   def get_int(self, value: int = None):
     return 0 if value is None else value
+
+  def typing_optional_get_int(self, value: Optional[int] = None):
+    return 0 if value is None else value
+
+  def typing_union_get_int(self, value: Union[int, str] = None):
+    return 0 if value is None else value