8000 Support generic partial types for attributes (#8044) · python/mypy@4e021d9 · GitHub
[go: up one dir, main page]
Skip to content

Commit 4e021d9

Browse files
ilevkivskyiilevkivskyi
authored
Support generic partial types for attributes (#8044)
Work towards #1055 Currently partial types are supported for local variables. However, only partial `None` types are supported for `self` attributes. This PR adds the same level of support to generic partial types. They follow mostly the same rules: * A partial type can be refined in the _same_ method where it is defined. * But a partial type from class body can not be refined in a method, as if `local_partial_types = True`. The logic is pretty simple: the `.node` attribute for `self.attr` expressions is set to `None`, so I added a little helper to get it from the class symbol table instead.
1 parent 3930bbf commit 4e021d9

File tree

4 files changed

+223
-4
lines changed

4 files changed

+223
-4
lines changed

‎mypy/checker.py

Lines changed: 20 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2153,13 +2153,20 @@ def try_infer_partial_generic_type_from_assignment(self,
21532153
if foo():
21542154
x = [1] # Infer List[int] as type of 'x'
21552155
"""
2156+
var = None
21562157
if (isinstance(lvalue, NameExpr)
21572158
and isinstance(lvalue.node, Var)
21582159
and isinstance(lvalue.node.type, PartialType)):
21592160
var = lvalue.node
2160-
typ = lvalue.node.type
2161+
elif isinstance(lvalue, MemberExpr):
2162+
var = self.expr_checker.get_partial_self_var(lvalue)
2163+
if var is not None:
2164+
typ = var.type
2165+
assert isinstance(typ, PartialType)
21612166
if typ.type is None:
21622167
return
2168+
# TODO: some logic here duplicates the None partial type counterpart
2169+
# inlined in check_assignment(), see # 8043.
21632170
partial_types = self.find_partial_types(var)
21642171
if partial_types is None:
21652172
return
@@ -2993,8 +3000,12 @@ def check_indexed_assignment(self, lvalue: IndexExpr,
29933000
def try_infer_partial_type_from_indexed_assignment(
29943001
self, lvalue: IndexExpr, rvalue: Expression) -> None:
29953002
# TODO: Should we share some of this with try_infer_partial_type?
3003+
var = None
29963004
if isinstance(lvalue.base, RefExpr) and isinstance(lvalue.base.node, Var):
29973005
var = lvalue.base.node
3006+
elif isinstance(lvalue.base, MemberExpr):
3007+
var = self.expr_checker.get_partial_self_var(lvalue.base)
3008+
if isinstance(var, Var):
29983009
if isinstance(var.type, PartialType):
29993010
type_type = var.type.type
30003011
if type_type is None:
@@ -4331,7 +4342,14 @@ def find_partial_types_in_all_scopes(
43314342
# All scopes within the outermost function are active. Scopes out of
43324343
# the outermost function are inactive to allow local reasoning (important
43334344
# for fine-grained incremental mode).
4334-
scope_active = (not self.options.local_partial_types
4345+
disallow_other_scopes = self.options.local_partial_types
4346+
4347+
if isinstance(var.type, PartialType) and var.type.type is not None and var.info:
4348+
# This is an ugly hack to make partial generic self attributes behave
4349+
# as if --local-partial-types is always on (because it used to be like this).
4350+
disallow_other_scopes = True
4351+
4352+
scope_active = (not disallow_other_scopes
43354353
or scope.is_local == self.partial_types[-1].is_local)
43364354
return scope_active, scope.is_local, scope.map
43374355
return False, False, None

‎mypy/checkexpr.py

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -537,6 +537,25 @@ def check_typeddict_call_with_kwargs(self, callee: TypedDictType,
537537

538538
return callee
539539

540+
def get_partial_self_var(self, expr: MemberExpr) -> Optional[Var]:
541+
"""Get variable node for a partial self attribute.
542+
543+
If the expression is not a self attribute, or attribute is not variable,
544+
or variable is not partial, return None.
545+
"""
546+
if not (isinstance(expr.expr, NameExpr) and
547+
isinstance(expr.expr.node, Var) and expr.expr.node.is_self):
548+
# Not a self.attr expression.
549+
return None
550+
info = self.chk.scope.enclosing_class()
551+
if not info or expr.name not in info.names:
552+
# Don't mess with partial types in superclasses.
553+
return None
554+
sym = info.names[expr.name]
555+
if isinstance(sym.node, Var) and isinstance(sym.node.type, PartialType):
556+
return sym.node
557+
return None
558+
540559
# Types and methods that can be used to infer partial types.
541560
item_args = {'builtins.list': ['append'],
542561
'builtins.set': ['add', 'discard'],
@@ -550,6 +569,8 @@ def check_typeddict_call_with_kwargs(self, callee: TypedDictType,
550569
def try_infer_partial_type(self, e: CallExpr) -> None:
551570
if isinstance(e.callee, MemberExpr) and isinstance(e.callee.expr, RefExpr):
552571
var = e.callee.expr.node
572+
if var is None and isinstance(e.callee.expr, MemberExpr):
573+
var = self.get_partial_self_var(e.callee.expr)
553574
if not isinstance(var, Var):
554575
return
555576
partial_types = self.chk.find_partial_types(var)

‎test-data/unit/check-inference.test

Lines changed: 117 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1458,9 +1458,9 @@ class A:
14581458
class A:
14591459
def f(self) -> None:
14601460
# Attributes aren't supported right now.
1461-
self.a = [] # E: Need type annotation for 'a' (hint: "a: List[<type>] = ...")
1461+
self.a = []
14621462
self.a.append(1)
1463-
self.a.append('')
1463+
self.a.append('') # E: Argument 1 to "append" of "list" has incompatible type "str"; expected "int"
14641464
[builtins fixtures/list.pyi]
14651465

14661466
[case testInferListInitializedToEmptyInClassBodyAndOverriden]
@@ -1585,6 +1585,121 @@ oo.update(d)
15851585
reveal_type(oo) # N: Revealed type is 'collections.OrderedDict[builtins.int*, builtins.str*]'
15861586
[builtins fixtures/dict.pyi]
15871587

1588+
[case testInferAttributeInitializedToEmptyAndAssigned]
1589+
class C:
1590+
def __init__(self) -> None:
1591+
self.a = []
1592+
if bool():
1593+
self.a = [1]
1594+
reveal_type(C().a) # N: Revealed type is 'builtins.list[builtins.int*]'
1595+
[builtins fixtures/list.pyi]
1596+
1597+
[case testInferAttributeInitializedToEmptyAndAppended]
1598+
class C:
1599+
def __init__(self) -> None:
1600+
self.a = []
1601+
if bool():
1602+
self.a.append(1)
1603+
reveal_type(C().a) # N: Revealed type is 'builtins.list[builtins.int]'
1604+
[builtins fixtures/list.pyi]
1605+
1606+
[case testInferAttributeInitializedToEmptyAndAssignedItem]
1607+
class C:
1608+
def __init__(self) -> None:
1609+
self.a = {}
1610+
if bool():
1611+
self.a[0] = 'yes'
1612+
reveal_type(C().a) # N: Revealed type is 'builtins.dict[builtins.int, builtins.str]'
1613+
[builtins fixtures/dict.pyi]
1614+
1615+
[case testInferAttributeInitializedToNoneAndAssigned]
1616+
# flags: --strict-optional
1617+
class C:
1618+
def __init__(self) -> None:
1619+
self.a = None
1620+
if bool():
1621+
self.a = 1
1622+
reveal_type(C().a) # N: Revealed type is 'Union[builtins.int, None]'
1623+
1624+
[case testInferAttributeInitializedToEmptyNonSelf]
1625+
class C:
1626+
def __init__(self) -> None:
1627+
self.a = [] # E: Need type annotation for 'a' (hint: "a: List[<type>] = ...")
1628+
if bool():
1629+
a = self
1630+
a.a = [1]
1631+
a.a.append(1)
1632+
reveal_type(C().a) # N: Revealed type is 'builtins.list[Any]'
1633+
[builtins fixtures/list.pyi]
1634+
1635+
[case testInferAttributeInitializedToEmptyAndAssignedOtherMethod]
1636+
class C:
1637+
def __init__(self) -> None:
1638+
self.a = [] # E: Need type annotation for 'a' (hint: "a: List[<type>] = ...")
1639+
def meth(self) -> None:
1640+
self.a = [1]
1641+
reveal_type(C().a) # N: Revealed type is 'builtins.list[Any]'
1642+
[builtins fixtures/list.pyi]
1643+
1644+
[case testInferAttributeInitializedToEmptyAndAppendedOtherMethod]
1645+
class C:
1646+
def __init__(self) -> None:
1647+
self.a = [] # E: Need type annotation for 'a' (hint: "a: List[<type>] = ...")
1648+
def meth(self) -> None:
1649+
self.a.append(1)
1650+
reveal_type(C().a) # N: Revealed type is 'builtins.list[Any]'
1651+
[builtins fixtures/list.pyi]
1652+
1653+
[case testInferAttributeInitializedToEmptyAndAssignedItemOtherMethod]
1654+
class C:
1655+
def __init__(self) -> None:
1656+
self.a = {} # E: Need type annotation for 'a' (hint: "a: Dict[<type>, <type>] = ...")
1657+
def meth(self) -> None:
1658+
self.a[0] = 'yes'
1659+
reveal_type(C().a) # N: Revealed type is 'builtins.dict[Any, Any]'
1660+
[builtins fixtures/dict.pyi]
1661+
1662+
[case testInferAttributeInitializedToNoneAndAssignedOtherMethod]
1663+
# flags: --strict-optional
1664+
class C:
1665+
def __init__(self) -> None:
1666+
self.a = None
1667+
def meth(self) -> None:
1668+
self.a = 1 # E: Incompatible types in assignment (expression has type "int", variable has type "None")
1669+
reveal_type(C().a) # N: Revealed type is 'None'
1670+
1671+
[case testInferAttributeInitializedToEmptyAndAssignedClassBody]
1672+
class C:
1673+
a = [] # E: Need type annotation for 'a' (hint: "a: List[<type>] = ...")
1674+
def __init__(self) -> None:
1675+
self.a = [1]
1676+
reveal_type(C().a) # N: Revealed type is 'builtins.list[Any]'
1677+
[builtins fixtures/list.pyi]
1678+
1679+
[case testInferAttributeInitializedToEmptyAndAppendedClassBody]
1680+
class C:
1681+
a = [] # E: Need type annotation for 'a' (hint: "a: List[<type>] = ...")
1682+
def __init__(self) -> None:
1683+
self.a.append(1)
1684+
reveal_type(C().a) # N: Revealed type is 'builtins.list[Any]'
1685+
[builtins fixtures/list.pyi]
1686+
1687+
[case testInferAttributeInitializedToEmptyAndAssignedItemClassBody]
1688+
class C:
1689+
a = {} # E: Need type annotation for 'a' (hint: "a: Dict[<type>, <type>] = ...")
1690+
def __init__(self) -> None:
1691+
self.a[0] = 'yes'
1692+
reveal_type(C().a) # N: Revealed type is 'builtins.dict[Any, Any]'
1693+
[builtins fixtures/dict.pyi]
1694+
1695+
[case testInferAttributeInitializedToNoneAndAssignedClassBody]
1696+
# flags: --strict-optional
1697+
class C:
1698+
a = None
1699+
def __init__(self) -> None:
1700+
self.a = 1
1701+
reveal_type(C().a) # N: Revealed type is 'Union[builtins.int, None]'
1702+
15881703

15891704
-- Inferring types of variables first initialized to None (partial types)
15901705
-- ----------------------------------------------------------------------

‎test-data/unit/fine-grained.test

Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2714,6 +2714,7 @@ class C:
27142714
class D:
27152715
def __init__(self) -> None:
27162716
self.x = {}
2717+
def meth(self) -> None:
27172718
self.x['a'] = 'b'
27182719
[file a.py]
27192720
def g() -> None: pass
@@ -2731,6 +2732,7 @@ class D:
27312732
def __init__(self) -> None:
27322733
a.g()
27332734
self.x = {}
2735+
def meth(self) -> None:
27342736
self.x['a'] = 'b'
27352737
[file a.py]
27362738
def g() -> None: pass
@@ -2742,6 +2744,69 @@ main:5: error: Need type annotation for 'x' (hint: "x: Dict[<type>, <type>] = ..
27422744
==
27432745
main:5: error: Need type annotation for 'x' (hint: "x: Dict[<type>, <type>] = ...")
27442746

2747+
[case testRefreshPartialTypeInferredAttributeIndex]
2748+
from c import C
2749+
reveal_type(C().a)
2750+
[file c.py]
2751+
from b import f
2752+
class C:
2753+
def __init__(self) -> None:
2754+
self.a = {}
2755+
if bool():
2756+
self.a[0] = f()
2757+
[file b.py]
2758+
def f() -> int: ...
2759+
[file b.py.2]
2760+
from typing import List
2761+
def f() -> str: ...
2762+
[builtins fixtures/dict.pyi]
2763+
[out]
2764+
main:2: note: Revealed type is 'builtins.dict[builtins.int, builtins.int]'
2765+
==
2766+
main:2: note: Revealed type is 'builtins.dict[builtins.int, builtins.str]'
2767+
2768+
[case testRefreshPartialTypeInferredAttributeAssign]
2769+
from c import C
2770+
reveal_type(C().a)
2771+
[file c.py]
2772+
from b import f
2773+
class C:
2774+
def __init__(self) -> None:
2775+
self.a = []
2776+
if bool():
2777+
self.a = f()
2778+
[file b.py]
2779+
from typing import List
2780+
def f() -> List[int]: ...
2781+
[file b.py.2]
2782+
from typing import List
2783+
def f() -> List[str]: ...
2784+
[builtins fixtures/list.pyi]
2785+
[out]
2786+
main:2: note: Revealed type is 'builtins.list[builtins.int]'
2787+
==
2788+
main:2: note: Revealed type is 'builtins.list[builtins.str]'
2789+
2790+
[case testRefreshPartialTypeInferredAttributeAppend]
2791+
from c import C
2792+
reveal_type(C().a)
2793+
[file c.py]
2794+
from b import f
2795+
class C:
2796+
def __init__(self) -> None:
2797+
self.a = []
2798+
if bool():
2799+
self.a.append(f())
2800+
[file b.py]
2801+
def f() -> int: ...
2802+
[file b.py.2]
2803+
def f() -> str: ...
2804+
[builtins fixtures/list.pyi]
2805+
[out]
2806+
main:2: note: Revealed type is 'builtins.list[builtins.int]'
2807+
==
2808+
main:2: note: Revealed type is 'builtins.list[builtins.str]'
2809+
27452810
[case testRefreshTryExcept]
27462811
import a
27472812
def f() -> None:

0 commit comments

Comments
 (0)
0