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

fire/helptext.py

@@ -35,7 +35,9 @@
 
 import collections
 import itertools
+import re
 import sys
+import typing
 
 from fire import completion
 from fire import custom_descriptions
@@ -538,12 +540,12 @@ def _GetArgType(arg, spec):
     arg_type = spec.annotations[arg]
     try:
       if sys.version_info[0:2] >= (3, 3):
+        if isinstance(arg_type, typing._GenericAlias):
+          arg_type = re.search(r'\[(.*?)\]', repr(arg_type)).group(1)
+          return arg_type
         return arg_type.__qualname__
       return arg_type.__name__
     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

@@ -164,6 +164,42 @@ def testHelpTextFunctionWithTypesAndDefaultNone(self):
         help_screen)
     self.assertNotIn('NOTES', help_screen)
 
+  @testutils.skipIf(
+      sys.version_info[0:2] < (3, 5),
+      'Python < 3.5 does not support type hints.')
+  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)
+
+  @testutils.skipIf(
+      sys.version_info[0:2] < (3, 5),
+      'Python < 3.5 does not support type hints.')
+  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[int, str]\n        Default: None',
+        help_screen)
+    self.assertNotIn('NOTES', help_screen)
+
   @testutils.skipIf(
       sys.version_info[0:2] < (3, 5),
       'Python < 3.5 does not support type hints.')

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