Keep NoneType in Union TypeVar values

Author: cdce8p

mypy/expandtype.py

@@ -50,25 +50,37 @@
 
 @overload
 def expand_type(
-    typ: ProperType, env: Mapping[TypeVarId, Type], allow_erased_callables: bool = ...
+    typ: ProperType,
+    env: Mapping[TypeVarId, Type],
+    allow_erased_callables: bool = ...,
+    *,
+    keep_none_type: bool = ...,
 ) -> ProperType:
     ...
 
 
 @overload
 def expand_type(
-    typ: Type, env: Mapping[TypeVarId, Type], allow_erased_callables: bool = ...
+    typ: Type,
+    env: Mapping[TypeVarId, Type],
+    allow_erased_callables: bool = ...,
+    *,
+    keep_none_type: bool = ...,
 ) -> Type:
     ...
 
 
 def expand_type(
-    typ: Type, env: Mapping[TypeVarId, Type], allow_erased_callables: bool = False
+    typ: Type,
+    env: Mapping[TypeVarId, Type],
+    allow_erased_callables: bool = False,
+    *,
+    keep_none_type: bool = False,
 ) -> Type:
     """Substitute any type variable references in a type given by a type
     environment.
     """
-    return typ.accept(ExpandTypeVisitor(env, allow_erased_callables))
+    return typ.accept(ExpandTypeVisitor(env, allow_erased_callables, keep_none_type))
 
 
 @overload
@@ -183,10 +195,14 @@ class ExpandTypeVisitor(TypeVisitor[Type]):
     variables: Mapping[TypeVarId, Type]  # TypeVar id -> TypeVar value
 
     def __init__(
-        self, variables: Mapping[TypeVarId, Type], allow_erased_callables: bool = False
+        self,
+        variables: Mapping[TypeVarId, Type],
+        allow_erased_callables: bool = False,
+        keep_none_type: bool = False,
     ) -> None:
         self.variables = variables
         self.allow_erased_callables = allow_erased_callables
+        self.keep_none_type = keep_none_type
         self.recursive_guard: set[Type | tuple[int, Type]] = set()
 
     def visit_unbound_type(self, t: UnboundType) -> Type:
@@ -470,7 +486,7 @@ def visit_union_type(self, t: UnionType) -> Type:
         # might be subtypes of others, however calling  make_simplified_union()
         # can cause recursion, so we just remove strict duplicates.
         return UnionType.make_union(
-            remove_trivial(flatten_nested_unions(expanded)), t.line, t.column
+            remove_trivial(flatten_nested_unions(expanded), self.keep_none_type), t.line, t.column
         )
 
     def visit_partial_type(self, t: PartialType) -> Type:

mypy/typeanal.py

@@ -1692,7 +1692,7 @@ def fix_instance(
             args.append(arg)
         env[tv.id] = arg
     t.args = tuple(args)
-    fixed = expand_type(t, env)
+    fixed = expand_type(t, env, keep_none_type=True)
     assert isinstance(fixed, Instance)
     t.args = fixed.args
 

mypy/types.py

@@ -3679,7 +3679,7 @@ def store_argument_type(
     defn.arguments[i].variable.type = arg_type
 
 
-def remove_trivial(types: Iterable[Type]) -> list[Type]:
+def remove_trivial(types: Iterable[Type], keep_none_type: bool = False) -> list[Type]:
     """Make trivial simplifications on a list of types without calling is_subtype().
 
     This makes following simplifications:
@@ -3694,7 +3694,7 @@ def remove_trivial(types: Iterable[Type]) -> list[Type]:
         p_t = get_proper_type(t)
         if isinstance(p_t, UninhabitedType):
             continue
-        if isinstance(p_t, NoneType) and not state.strict_optional:
+        if isinstance(p_t, NoneType) and not state.strict_optional and not keep_none_type:
             removed_none = True
             continue
         if isinstance(p_t, Instance) and p_t.type.fullname == "builtins.object":