Improve error message for partial None with --local-partial-types (#12822)

When --local-partial-types is set and we can't infer a complete type for
a type that we initially inferred as partial None, show an error message
that suggests to add a type annotation of the form Optional[<type>].

Co-authored-by: hauntsaninja <[email protected]>

Author: pranavrajpal

mypy/messages.py

@@ -1526,23 +1526,30 @@ def need_annotation_for_var(
     ) -> None:
         hint = ""
         has_variable_annotations = not python_version or python_version >= (3, 6)
+        pep604_supported = not python_version or python_version >= (3, 10)
+        # type to recommend the user adds
+        recommended_type = None
         # Only gives hint if it's a variable declaration and the partial type is a builtin type
-        if (
-            python_version
-            and isinstance(node, Var)
-            and isinstance(node.type, PartialType)
-            and node.type.type
-            and node.type.type.fullname in reverse_builtin_aliases
-        ):
-            alias = reverse_builtin_aliases[node.type.type.fullname]
-            alias = alias.split(".")[-1]
+        if python_version and isinstance(node, Var) and isinstance(node.type, PartialType):
             type_dec = "<type>"
-            if alias == "Dict":
-                type_dec = f"{type_dec}, {type_dec}"
+            if not node.type.type:
+                # partial None
+                if pep604_supported:
+                    recommended_type = f"{type_dec} | None"
+                else:
+                    recommended_type = f"Optional[{type_dec}]"
+            elif node.type.type.fullname in reverse_builtin_aliases:
+                # partial types other than partial None
+                alias = reverse_builtin_aliases[node.type.type.fullname]
+                alias = alias.split(".")[-1]
+                if alias == "Dict":
+                    type_dec = f"{type_dec}, {type_dec}"
+                recommended_type = f"{alias}[{type_dec}]"
+        if recommended_type is not None:
             if has_variable_annotations:
-                hint = f' (hint: "{node.name}: {alias}[{type_dec}] = ...")'
+                hint = f' (hint: "{node.name}: {recommended_type} = ...")'
             else:
-                hint = f' (hint: "{node.name} = ...  # type: {alias}[{type_dec}]")'
+                hint = f' (hint: "{node.name} = ...  # type: {recommended_type}")'
 
         if has_variable_annotations:
             needed = "annotation"

test-data/unit/check-inference.test

@@ -2393,7 +2393,7 @@ if bool():
 
 [case testLocalPartialTypesWithGlobalInitializedToNone]
 # flags: --local-partial-types
-x = None # E: Need type annotation for "x"
+x = None # E: Need type annotation for "x" (hint: "x: Optional[<type>] = ...")
 
 def f() -> None:
     global x
@@ -2404,7 +2404,7 @@ reveal_type(x) # N: Revealed type is "None"
 
 [case testLocalPartialTypesWithGlobalInitializedToNone2]
 # flags: --local-partial-types
-x = None # E: Need type annotation for "x"
+x = None # E: Need type annotation for "x" (hint: "x: Optional[<type>] = ...")
 
 def f():
     global x
@@ -2453,7 +2453,7 @@ reveal_type(a)  # N: Revealed type is "builtins.str"
 [case testLocalPartialTypesWithClassAttributeInitializedToNone]
 # flags: --local-partial-types
 class A:
-    x = None # E: Need type annotation for "x"
+    x = None # E: Need type annotation for "x" (hint: "x: Optional[<type>] = ...")
 
     def f(self) -> None:
         self.x = 1
@@ -2636,7 +2636,7 @@ from typing import List
 def f(x): pass
 
 class A:
-    x = None  # E: Need type annotation for "x"
+    x = None  # E: Need type annotation for "x" (hint: "x: Optional[<type>] = ...")
 
     def f(self, p: List[str]) -> None:
         self.x = f(p)
@@ -2646,15 +2646,15 @@ class A:
 [case testLocalPartialTypesAccessPartialNoneAttribute]
 # flags: --local-partial-types
 class C:
-    a = None  # E: Need type annotation for "a"
+    a = None  # E: Need type annotation for "a" (hint: "a: Optional[<type>] = ...")
 
     def f(self, x) -> None:
         C.a.y  # E: Item "None" of "Optional[Any]" has no attribute "y"
 
 [case testLocalPartialTypesAccessPartialNoneAttribute2]
 # flags: --local-partial-types
 class C:
-    a = None  # E: Need type annotation for "a"
+    a = None  # E: Need type annotation for "a" (hint: "a: Optional[<type>] = ...")
 
     def f(self, x) -> None:
         self.a.y  # E: Item "None" of "Optional[Any]" has no attribute "y"
@@ -3248,6 +3248,10 @@ if x:
     reveal_type(x)  # N: Revealed type is "builtins.bytes"
 [builtins fixtures/dict.pyi]
 
+[case testSuggestPep604AnnotationForPartialNone]
+# flags: --local-partial-types --python-version 3.10
+x = None # E: Need type annotation for "x" (hint: "x: <type> | None = ...")
+
 [case testTupleContextFromIterable]
 from typing import TypeVar, Iterable, List, Union
 

test-data/unit/fine-grained.test

@@ -4305,9 +4305,9 @@ y = 0
 [file a.py.2]
 y = ''
 [out]
-main:4: error: Need type annotation for "x"
+main:4: error: Need type annotation for "x" (hint: "x: Optional[<type>] = ...")
 ==
-main:4: error: Need type annotation for "x"
+main:4: error: Need type annotation for "x" (hint: "x: Optional[<type>] = ...")
 
 [case testNonePartialType2]
 import a
@@ -4323,9 +4323,9 @@ y = 0
 [file a.py.2]
 y = ''
 [out]
-main:4: error: Need type annotation for "x"
+main:4: error: Need type annotation for "x" (hint: "x: Optional[<type>] = ...")
 ==
-main:4: error: Need type annotation for "x"
+main:4: error: Need type annotation for "x" (hint: "x: Optional[<type>] = ...")
 
 [case testNonePartialType3]
 import a
@@ -4337,7 +4337,7 @@ def f() -> None:
     y = ''
 [out]
 ==
-a.py:1: error: Need type annotation for "y"
+a.py:1: error: Need type annotation for "y" (hint: "y: Optional[<type>] = ...")
 
 [case testNonePartialType4]
 import a
@@ -4353,7 +4353,7 @@ def f() -> None:
     global y
     y = ''
 [out]
-a.py:1: error: Need type annotation for "y"
+a.py:1: error: Need type annotation for "y" (hint: "y: Optional[<type>] = ...")
 ==
 
 [case testSkippedClass1]