Fix UnionType handling as with Union for Optional values

jonaslb · jonaslb · commit 1f9dd82290ca · 2023-02-04T03:10:24.000+01:00
diff --git a/tests/test_type_conversion.py b/tests/test_type_conversion.py
@@ -1,3 +1,4 @@
+import sys
 from enum import Enum
 from pathlib import Path
 from typing import List, Optional, Tuple
@@ -28,6 +29,28 @@ def opt(user: Optional[str] = None):
     assert "User: Camila" in result.output
 
 
+@pytest.mark.skipif(
+    sys.version_info < (3, 10), reason="The | operator for types was new in 3.10"
+)
+def test_union_type_optional():
+    app = typer.Typer()
+
+    @app.command()
+    def opt(user: str | None = None):
+        if user:
+            print(f"User: {user}")
+        else:
+            print("No user")
+
+    result = runner.invoke(app)
+    assert result.exit_code == 0
+    assert "No user" in result.output
+
+    result = runner.invoke(app, ["--user", "Camila"])
+    assert result.exit_code == 0
+    assert "User: Camila" in result.output
+
+
 def test_no_type():
     app = typer.Typer()
 
diff --git a/typer/_compat_utils.py b/typer/_compat_utils.py
@@ -1,5 +1,38 @@
+import sys
+from typing import Union
+
 import click
 
+if sys.version_info >= (3, 8):
+    from typing import get_args as _get_args
+    from typing import get_origin as _get_origin
+elif sys.version_info >= (3, 7):
+    from typing_extensions import get_args as _get_args
+    from typing_extensions import get_origin as _get_origin
+else:
+    # These methods do not handle all the same details as the imported ones.
+    # However on Python 3.6 they should be sufficient.
+    # typer <= 0.7.0 used this implementation.
+
+    def _get_origin(arg):
+        return getattr(arg, "__origin__", None)
+
+    def _get_args(arg):
+        return getattr(arg, "__args__", None)
+
+
+# Assigning variables to mark them as exported with mypy
+get_origin = _get_origin
+get_args = _get_args
+
+if sys.version_info >= (3, 10):
+    # Type ignore since mypy does not think UnionType exists
+    from types import UnionType  # type: ignore
+
+    UNION_TYPES = (Union, UnionType)
+else:
+    UNION_TYPES = (Union,)
+
 
 def _get_click_major() -> int:
     return int(click.__version__.split(".")[0])
diff --git a/typer/main.py b/typer/main.py
@@ -13,6 +13,7 @@
 
 import click
 
+from ._compat_utils import UNION_TYPES, get_args, get_origin
 from .completion import get_completion_inspect_parameters
 from .core import MarkupMode, TyperArgument, TyperCommand, TyperGroup, TyperOption
 from .models import (
@@ -810,30 +811,31 @@ def get_click_param(
     is_tuple = False
     parameter_type: Any = None
     is_flag = None
-    origin = getattr(main_type, "__origin__", None)
+    origin = get_origin(main_type)
+
     if origin is not None:
-        # Handle Optional[SomeType]
-        if origin is Union:
+        # Handle SomeType | None and Optional[SomeType]
+        if origin in UNION_TYPES:
             types = []
-            for type_ in main_type.__args__:
+            for type_ in get_args(main_type):
                 if type_ is NoneType:
                     continue
                 types.append(type_)
             assert len(types) == 1, "Typer Currently doesn't support Union types"
             main_type = types[0]
-            origin = getattr(main_type, "__origin__", None)
+            origin = get_origin(main_type)
         # Handle Tuples and Lists
         if lenient_issubclass(origin, List):
-            main_type = main_type.__args__[0]
-            assert not getattr(
-                main_type, "__origin__", None
+            main_type = get_args(main_type)[0]
+            assert not get_origin(
+                main_type
             ), "List types with complex sub-types are not currently supported"
             is_list = True
         elif lenient_issubclass(origin, Tuple):  # type: ignore
             types = []
-            for type_ in main_type.__args__:
-                assert not getattr(
-                    type_, "__origin__", None
+            for type_ in get_args(main_type):
+                assert not get_origin(
+                    type_
                 ), "Tuple types with complex sub-types are not currently supported"
                 types.append(
                     get_click_type(annotation=type_, parameter_info=parameter_info)
@@ -848,7 +850,7 @@ def get_click_param(
     if is_list:
         convertor = generate_list_convertor(convertor)
     if is_tuple:
-        convertor = generate_tuple_convertor(main_type.__args__)
+        convertor = generate_tuple_convertor(get_args(main_type))
     if isinstance(parameter_info, OptionInfo):
         if main_type is bool and not (parameter_info.is_flag is False):
             is_flag = True
@@ -1002,7 +1004,7 @@ def get_param_completion(
     incomplete_name = None
     unassigned_params = [param for param in parameters.values()]
     for param_sig in unassigned_params[:]:
-        origin = getattr(param_sig.annotation, "__origin__", None)
+        origin = get_origin(param_sig.annotation)
         if lenient_issubclass(param_sig.annotation, click.Context):
             ctx_name = param_sig.name
             unassigned_params.remove(param_sig)
