Use supertype context for variable type inference (#13494)

Author: ilevkivskyi

mypy/checker.py

@@ -1624,6 +1624,8 @@ def check_slots_definition(self, typ: Type, context: Context) -> None:
 
     def check_match_args(self, var: Var, typ: Type, context: Context) -> None:
         """Check that __match_args__ contains literal strings"""
+        if not self.scope.active_class():
+            return
         typ = get_proper_type(typ)
         if not isinstance(typ, TupleType) or not all(
             [is_string_literal(item) for item in typ.items]
@@ -2686,7 +2688,8 @@ def check_assignment(
                 self.check_indexed_assignment(index_lvalue, rvalue, lvalue)
 
             if inferred:
-                rvalue_type = self.expr_checker.accept(rvalue)
+                type_context = self.get_variable_type_context(inferred)
+                rvalue_type = self.expr_checker.accept(rvalue, type_context=type_context)
                 if not (
                     inferred.is_final
                     or (isinstance(lvalue, NameExpr) and lvalue.name == "__match_args__")
@@ -2698,6 +2701,27 @@ def check_assignment(
     # (type, operator) tuples for augmented assignments supported with partial types
     partial_type_augmented_ops: Final = {("builtins.list", "+"), ("builtins.set", "|")}
 
+    def get_variable_type_context(self, inferred: Var) -> Type | None:
+        type_contexts = []
+        if inferred.info:
+            for base in inferred.info.mro[1:]:
+                base_type, base_node = self.lvalue_type_from_base(inferred, base)
+                if base_type and not (
+                    isinstance(base_node, Var) and base_node.invalid_partial_type
+                ):
+                    type_contexts.append(base_type)
+        # Use most derived supertype as type context if available.
+        if not type_contexts:
+            return None
+        candidate = type_contexts[0]
+        for other in type_contexts:
+            if is_proper_subtype(other, candidate):
+                candidate = other
+            elif not is_subtype(candidate, other):
+                # Multiple incompatible candidates, cannot use any of them as context.
+                return None
+        return candidate
+
     def try_infer_partial_generic_type_from_assignment(
         self, lvalue: Lvalue, rvalue: Expression, op: str
     ) -> None:
@@ -5870,7 +5894,9 @@ def enter_partial_types(
                         self.msg.need_annotation_for_var(var, context, self.options.python_version)
                         self.partial_reported.add(var)
                     if var.type:
-                        var.type = self.fixup_partial_type(var.type)
+                        fixed = self.fixup_partial_type(var.type)
+                        var.invalid_partial_type = fixed != var.type
+                        var.type = fixed
 
     def handle_partial_var_type(
         self, typ: PartialType, is_lvalue: bool, node: Var, context: Context

mypy/nodes.py

@@ -939,6 +939,7 @@ def deserialize(cls, data: JsonDict) -> Decorator:
     "explicit_self_type",
     "is_ready",
     "is_inferred",
+    "invalid_partial_type",
     "from_module_getattr",
     "has_explicit_value",
     "allow_incompatible_override",
@@ -975,6 +976,7 @@ class Var(SymbolNode):
         "from_module_getattr",
         "has_explicit_value",
         "allow_incompatible_override",
+        "invalid_partial_type",
     )
 
     def __init__(self, name: str, type: mypy.types.Type | None = None) -> None:
@@ -1024,6 +1026,9 @@ def __init__(self, name: str, type: mypy.types.Type | None = None) -> None:
         self.has_explicit_value = False
         # If True, subclasses can override this with an incompatible type.
         self.allow_incompatible_override = False
+        # If True, this means we didn't manage to infer full type and fall back to
+        # something like list[Any]. We may decide to not use such types as context.
+        self.invalid_partial_type = False
 
     @property
     def name(self) -> str:

mypy/semanal.py

@@ -3007,7 +3007,10 @@ def process_type_annotation(self, s: AssignmentStmt) -> None:
             ):
                 self.fail("All protocol members must have explicitly declared types", s)
             # Set the type if the rvalue is a simple literal (even if the above error occurred).
-            if len(s.lvalues) == 1 and isinstance(s.lvalues[0], RefExpr):
+            # We skip this step for type scope because it messes up with class attribute
+            # inference for literal types (also annotated and non-annotated variables at class
+            # scope are semantically different, so we should not souch statement type).
+            if len(s.lvalues) == 1 and isinstance(s.lvalues[0], RefExpr) and not self.type:
                 if s.lvalues[0].is_inferred_def:
                     s.type = self.analyze_simple_literal_type(s.rvalue, s.is_final_def)
         if s.type:
@@ -3026,7 +3029,6 @@ def is_annotated_protocol_member(self, s: AssignmentStmt) -> bool:
 
     def analyze_simple_literal_type(self, rvalue: Expression, is_final: bool) -> Type | None:
         """Return builtins.int if rvalue is an int literal, etc.
-
         If this is a 'Final' context, we return "Literal[...]" instead."""
         if self.options.semantic_analysis_only or self.function_stack:
             # Skip this if we're only doing the semantic analysis pass.

test-data/unit/check-classes.test

@@ -4317,7 +4317,7 @@ class C(B):
     x = object()
 [out]
 main:4: error: Incompatible types in assignment (expression has type "str", base class "A" defined the type as "int")
-main:6: error: Incompatible types in assignment (expression has type "object", base class "B" defined the type as "str")
+main:6: error: Incompatible types in assignment (expression has type "object", base class "A" defined the type as "int")
 
 [case testClassOneErrorPerLine]
 class A:
@@ -4327,15 +4327,15 @@ class B(A):
   x = 1.0
 [out]
 main:4: error: Incompatible types in assignment (expression has type "str", base class "A" defined the type as "int")
-main:5: error: Incompatible types in assignment (expression has type "str", base class "A" defined the type as "int")
+main:5: error: Incompatible types in assignment (expression has type "float", base class "A" defined the type as "int")
 
 [case testClassIgnoreType_RedefinedAttributeAndGrandparentAttributeTypesNotIgnored]
 class A:
     x = 0
 class B(A):
     x = ''  # type: ignore
 class C(B):
-    x = ''
+    x = ''  # E: Incompatible types in assignment (expression has type "str", base class "A" defined the type as "int")
 [out]
 
 [case testClassIgnoreType_RedefinedAttributeTypeIgnoredInChildren]

test-data/unit/check-inference.test

@@ -3263,3 +3263,68 @@ from typing import Dict, Iterable, Tuple, Union
 def foo(x: Union[Tuple[str, Dict[str, int], str], Iterable[object]]) -> None: ...
 foo(("a", {"a": "b"}, "b"))
 [builtins fixtures/dict.pyi]
+
+[case testUseSupertypeAsInferenceContext]
+# flags: --strict-optional
+from typing import List, Optional
+
+class B:
+    x: List[Optional[int]]
+
+class C(B):
+    x = [1]
+
+reveal_type(C().x)  # N: Revealed type is "builtins.list[Union[builtins.int, None]]"
+[builtins fixtures/list.pyi]
+
+[case testUseSupertypeAsInferenceContextInvalidType]
+from typing import List
+class P:
+    x: List[int]
+class C(P):
+    x = ['a']  # E: List item 0 has incompatible type "str"; expected "int"
+[builtins fixtures/list.pyi]
+
+[case testUseSupertypeAsInferenceContextPartial]
+from typing import List
+
+class A:
+    x: List[str]
+
+class B(A):
+    x = []
+
+reveal_type(B().x)  # N: Revealed type is "builtins.list[builtins.str]"
+[builtins fixtures/list.pyi]
+
+[case testUseSupertypeAsInferenceContextPartialError]
+class A:
+    x = ['a', 'b']
+
+class B(A):
+    x = []
+    x.append(2)  # E: Argument 1 to "append" of "list" has incompatible type "int"; expected "str"
+[builtins fixtures/list.pyi]
+
+[case testUseSupertypeAsInferenceContextPartialErrorProperty]
+from typing import List
+
+class P:
+    @property
+    def x(self) -> List[int]: ...
+class C(P):
+    x = []
+
+C.x.append("no")  # E: Argument 1 to "append" of "list" has incompatible type "str"; expected "int"
+[builtins fixtures/list.pyi]
+
+[case testUseSupertypeAsInferenceContextConflict]
+from typing import List
+class P:
+    x: List[int]
+class M:
+    x: List[str]
+class C(P, M):
+    x = []  # E: Need type annotation for "x" (hint: "x: List[<type>] = ...")
+reveal_type(C.x)  # N: Revealed type is "builtins.list[Any]"
+[builtins fixtures/list.pyi]

test-data/unit/check-literal.test

@@ -2918,3 +2918,44 @@ def incorrect_return2() -> Union[Tuple[Literal[True], int], Tuple[Literal[False]
     else:
         return (bool(), 'oops')  # E: Incompatible return value type (got "Tuple[bool, str]", expected "Union[Tuple[Literal[True], int], Tuple[Literal[False], str]]")
 [builtins fixtures/bool.pyi]
+
+[case testLiteralSubtypeContext]
+from typing_extensions import Literal
+
+class A:
+    foo: Literal['bar', 'spam']
+class B(A):
+    foo = 'spam'
+
+reveal_type(B().foo)  # N: Revealed type is "Literal['spam']"
+[builtins fixtures/tuple.pyi]
+
+[case testLiteralSubtypeContextNested]
+from typing import List
+from typing_extensions import Literal
+
+class A:
+    foo: List[Literal['bar', 'spam']]
+class B(A):
+    foo = ['spam']
+
+reveal_type(B().foo)  # N: Revealed type is "builtins.list[Union[Literal['bar'], Literal['spam']]]"
+[builtins fixtures/tuple.pyi]
+
+[case testLiteralSubtypeContextGeneric]
+from typing_extensions import Literal
+from typing import Generic, List, TypeVar
+
+T = TypeVar("T", bound=str)
+
+class B(Generic[T]):
+    collection: List[T]
+    word: T
+
+class C(B[Literal["word"]]):
+   collection = ["word"]
+   word = "word"
+
+reveal_type(C().collection)  # N: Revealed type is "builtins.list[Literal['word']]"
+reveal_type(C().word)  # N: Revealed type is "Literal['word']"
+[builtins fixtures/tuple.pyi]