diff --git a/mypy/checker.py b/mypy/checker.py index 60625f6a3a26..d244b5255c30 100644 --- a/mypy/checker.py +++ b/mypy/checker.py @@ -2153,13 +2153,20 @@ def try_infer_partial_generic_type_from_assignment(self, if foo(): x = [1] # Infer List[int] as type of 'x' """ + var = None if (isinstance(lvalue, NameExpr) and isinstance(lvalue.node, Var) and isinstance(lvalue.node.type, PartialType)): var = lvalue.node - typ = lvalue.node.type + elif isinstance(lvalue, MemberExpr): + var = self.expr_checker.get_partial_self_var(lvalue) + if var is not None: + typ = var.type + assert isinstance(typ, PartialType) if typ.type is None: return + # TODO: some logic here duplicates the None partial type counterpart + # inlined in check_assignment(), see # 8043. partial_types = self.find_partial_types(var) if partial_types is None: return @@ -2993,8 +3000,12 @@ def check_indexed_assignment(self, lvalue: IndexExpr, def try_infer_partial_type_from_indexed_assignment( self, lvalue: IndexExpr, rvalue: Expression) -> None: # TODO: Should we share some of this with try_infer_partial_type? + var = None if isinstance(lvalue.base, RefExpr) and isinstance(lvalue.base.node, Var): var = lvalue.base.node + elif isinstance(lvalue.base, MemberExpr): + var = self.expr_checker.get_partial_self_var(lvalue.base) + if isinstance(var, Var): if isinstance(var.type, PartialType): type_type = var.type.type if type_type is None: @@ -4329,7 +4340,14 @@ def find_partial_types_in_all_scopes( # All scopes within the outermost function are active. Scopes out of # the outermost function are inactive to allow local reasoning (important # for fine-grained incremental mode). - scope_active = (not self.options.local_partial_types + disallow_other_scopes = self.options.local_partial_types + + if isinstance(var.type, PartialType) and var.type.type is not None and var.info: + # This is an ugly hack to make partial generic self attributes behave + # as if --local-partial-types is always on (because it used to be like this). + disallow_other_scopes = True + + scope_active = (not disallow_other_scopes or scope.is_local == self.partial_types[-1].is_local) return scope_active, scope.is_local, scope.map return False, False, None diff --git a/mypy/checkexpr.py b/mypy/checkexpr.py index 82bd86fba222..bb0c34b27ee8 100644 --- a/mypy/checkexpr.py +++ b/mypy/checkexpr.py @@ -537,6 +537,25 @@ def check_typeddict_call_with_kwargs(self, callee: TypedDictType, return callee + def get_partial_self_var(self, expr: MemberExpr) -> Optional[Var]: + """Get variable node for a partial self attribute. + + If the expression is not a self attribute, or attribute is not variable, + or variable is not partial, return None. + """ + if not (isinstance(expr.expr, NameExpr) and + isinstance(expr.expr.node, Var) and expr.expr.node.is_self): + # Not a self.attr expression. + return None + info = self.chk.scope.enclosing_class() + if not info or expr.name not in info.names: + # Don't mess with partial types in superclasses. + return None + sym = info.names[expr.name] + if isinstance(sym.node, Var) and isinstance(sym.node.type, PartialType): + return sym.node + return None + # Types and methods that can be used to infer partial types. item_args = {'builtins.list': ['append'], 'builtins.set': ['add', 'discard'], @@ -550,6 +569,8 @@ def check_typeddict_call_with_kwargs(self, callee: TypedDictType, def try_infer_partial_type(self, e: CallExpr) -> None: if isinstance(e.callee, MemberExpr) and isinstance(e.callee.expr, RefExpr): var = e.callee.expr.node + if var is None and isinstance(e.callee.expr, MemberExpr): + var = self.get_partial_self_var(e.callee.expr) if not isinstance(var, Var): return partial_types = self.chk.find_partial_types(var) diff --git a/test-data/unit/check-inference.test b/test-data/unit/check-inference.test index 8b5433fcf795..b2050a127ba6 100644 --- a/test-data/unit/check-inference.test +++ b/test-data/unit/check-inference.test @@ -1458,9 +1458,9 @@ class A: class A: def f(self) -> None: # Attributes aren't supported right now. - self.a = [] # E: Need type annotation for 'a' (hint: "a: List[] = ...") + self.a = [] self.a.append(1) - self.a.append('') + self.a.append('') # E: Argument 1 to "append" of "list" has incompatible type "str"; expected "int" [builtins fixtures/list.pyi] [case testInferListInitializedToEmptyInClassBodyAndOverriden] @@ -1585,6 +1585,121 @@ oo.update(d) reveal_type(oo) # N: Revealed type is 'collections.OrderedDict[builtins.int*, builtins.str*]' [builtins fixtures/dict.pyi] +[case testInferAttributeInitializedToEmptyAndAssigned] +class C: + def __init__(self) -> None: + self.a = [] + if bool(): + self.a = [1] +reveal_type(C().a) # N: Revealed type is 'builtins.list[builtins.int*]' +[builtins fixtures/list.pyi] + +[case testInferAttributeInitializedToEmptyAndAppended] +class C: + def __init__(self) -> None: + self.a = [] + if bool(): + self.a.append(1) +reveal_type(C().a) # N: Revealed type is 'builtins.list[builtins.int]' +[builtins fixtures/list.pyi] + +[case testInferAttributeInitializedToEmptyAndAssignedItem] +class C: + def __init__(self) -> None: + self.a = {} + if bool(): + self.a[0] = 'yes' +reveal_type(C().a) # N: Revealed type is 'builtins.dict[builtins.int, builtins.str]' +[builtins fixtures/dict.pyi] + +[case testInferAttributeInitializedToNoneAndAssigned] +# flags: --strict-optional +class C: + def __init__(self) -> None: + self.a = None + if bool(): + self.a = 1 +reveal_type(C().a) # N: Revealed type is 'Union[builtins.int, None]' + +[case testInferAttributeInitializedToEmptyNonSelf] +class C: + def __init__(self) -> None: + self.a = [] # E: Need type annotation for 'a' (hint: "a: List[] = ...") + if bool(): + a = self + a.a = [1] + a.a.append(1) +reveal_type(C().a) # N: Revealed type is 'builtins.list[Any]' +[builtins fixtures/list.pyi] + +[case testInferAttributeInitializedToEmptyAndAssignedOtherMethod] +class C: + def __init__(self) -> None: + self.a = [] # E: Need type annotation for 'a' (hint: "a: List[] = ...") + def meth(self) -> None: + self.a = [1] +reveal_type(C().a) # N: Revealed type is 'builtins.list[Any]' +[builtins fixtures/list.pyi] + +[case testInferAttributeInitializedToEmptyAndAppendedOtherMethod] +class C: + def __init__(self) -> None: + self.a = [] # E: Need type annotation for 'a' (hint: "a: List[] = ...") + def meth(self) -> None: + self.a.append(1) +reveal_type(C().a) # N: Revealed type is 'builtins.list[Any]' +[builtins fixtures/list.pyi] + +[case testInferAttributeInitializedToEmptyAndAssignedItemOtherMethod] +class C: + def __init__(self) -> None: + self.a = {} # E: Need type annotation for 'a' (hint: "a: Dict[, ] = ...") + def meth(self) -> None: + self.a[0] = 'yes' +reveal_type(C().a) # N: Revealed type is 'builtins.dict[Any, Any]' +[builtins fixtures/dict.pyi] + +[case testInferAttributeInitializedToNoneAndAssignedOtherMethod] +# flags: --strict-optional +class C: + def __init__(self) -> None: + self.a = None + def meth(self) -> None: + self.a = 1 # E: Incompatible types in assignment (expression has type "int", variable has type "None") +reveal_type(C().a) # N: Revealed type is 'None' + +[case testInferAttributeInitializedToEmptyAndAssignedClassBody] +class C: + a = [] # E: Need type annotation for 'a' (hint: "a: List[] = ...") + def __init__(self) -> None: + self.a = [1] +reveal_type(C().a) # N: Revealed type is 'builtins.list[Any]' +[builtins fixtures/list.pyi] + +[case testInferAttributeInitializedToEmptyAndAppendedClassBody] +class C: + a = [] # E: Need type annotation for 'a' (hint: "a: List[] = ...") + def __init__(self) -> None: + self.a.append(1) +reveal_type(C().a) # N: Revealed type is 'builtins.list[Any]' +[builtins fixtures/list.pyi] + +[case testInferAttributeInitializedToEmptyAndAssignedItemClassBody] +class C: + a = {} # E: Need type annotation for 'a' (hint: "a: Dict[, ] = ...") + def __init__(self) -> None: + self.a[0] = 'yes' +reveal_type(C().a) # N: Revealed type is 'builtins.dict[Any, Any]' +[builtins fixtures/dict.pyi] + +[case testInferAttributeInitializedToNoneAndAssignedClassBody] +# flags: --strict-optional +class C: + a = None + def __init__(self) -> None: + self.a = 1 +reveal_type(C().a) # N: Revealed type is 'Union[builtins.int, None]' + -- Inferring types of variables first initialized to None (partial types) -- ---------------------------------------------------------------------- diff --git a/test-data/unit/fine-grained.test b/test-data/unit/fine-grained.test index 0ee9d3094f47..11e83f560eee 100644 --- a/test-data/unit/fine-grained.test +++ b/test-data/unit/fine-grained.test @@ -2714,6 +2714,7 @@ class C: class D: def __init__(self) -> None: self.x = {} + def meth(self) -> None: self.x['a'] = 'b' [file a.py] def g() -> None: pass @@ -2731,6 +2732,7 @@ class D: def __init__(self) -> None: a.g() self.x = {} + def meth(self) -> None: self.x['a'] = 'b' [file a.py] def g() -> None: pass @@ -2742,6 +2744,69 @@ main:5: error: Need type annotation for 'x' (hint: "x: Dict[, ] = .. == main:5: error: Need type annotation for 'x' (hint: "x: Dict[, ] = ...") +[case testRefreshPartialTypeInferredAttributeIndex] +from c import C +reveal_type(C().a) +[file c.py] +from b import f +class C: + def __init__(self) -> None: + self.a = {} + if bool(): + self.a[0] = f() +[file b.py] +def f() -> int: ... +[file b.py.2] +from typing import List +def f() -> str: ... +[builtins fixtures/dict.pyi] +[out] +main:2: note: Revealed type is 'builtins.dict[builtins.int, builtins.int]' +== +main:2: note: Revealed type is 'builtins.dict[builtins.int, builtins.str]' + +[case testRefreshPartialTypeInferredAttributeAssign] +from c import C +reveal_type(C().a) +[file c.py] +from b import f +class C: + def __init__(self) -> None: + self.a = [] + if bool(): + self.a = f() +[file b.py] +from typing import List +def f() -> List[int]: ... +[file b.py.2] +from typing import List +def f() -> List[str]: ... +[builtins fixtures/list.pyi] +[out] +main:2: note: Revealed type is 'builtins.list[builtins.int]' +== +main:2: note: Revealed type is 'builtins.list[builtins.str]' + +[case testRefreshPartialTypeInferredAttributeAppend] +from c import C +reveal_type(C().a) +[file c.py] +from b import f +class C: + def __init__(self) -> None: + self.a = [] + if bool(): + self.a.append(f()) +[file b.py] +def f() -> int: ... +[file b.py.2] +def f() -> str: ... +[builtins fixtures/list.pyi] +[out] +main:2: note: Revealed type is 'builtins.list[builtins.int]' +== +main:2: note: Revealed type is 'builtins.list[builtins.str]' + [case testRefreshTryExcept] import a def f() -> None: