diff --git a/fire/helptext.py b/fire/helptext.py index 318d6276..27230ac3 100644 --- a/fire/helptext.py +++ b/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 '' diff --git a/fire/helptext_test.py b/fire/helptext_test.py index aeff5240..bc54da75 100644 --- a/fire/helptext_test.py +++ b/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 ', 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 ', 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( diff --git a/fire/test_components_py3.py b/fire/test_components_py3.py index 192302d3..b75ecaca 100644 --- a/fire/test_components_py3.py +++ b/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 \ No newline at end of file