Skip to content

Commit 9edc38d

Browse files
Michael0x2ailevkivskyi
authored andcommitted
Add support for using intelligent indexing with Final-declared variables (#6176)
This pull request adds support for intelligent indexing with variables declared using an implicit Final. This makes the following code work as expected: ``` my_tuple: Tuple[int, str, bool] idx: Final = 1 reveal_type(my_tuple[idx]) # Revealed type should be 'str' ```
1 parent bf0d795 commit 9edc38d

File tree

3 files changed

+48
-1
lines changed

3 files changed

+48
-1
lines changed

mypy/checkexpr.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2459,6 +2459,8 @@ def _get_value(self, index: Expression) -> Optional[int]:
24592459
if isinstance(operand, IntExpr):
24602460
return -1 * operand.value
24612461
typ = self.accept(index)
2462+
if isinstance(typ, Instance) and typ.final_value is not None:
2463+
typ = typ.final_value
24622464
if isinstance(typ, LiteralType) and isinstance(typ.value, int):
24632465
return typ.value
24642466
return None
@@ -2468,6 +2470,9 @@ def visit_typeddict_index_expr(self, td_type: TypedDictType, index: Expression)
24682470
item_name = index.value
24692471
else:
24702472
typ = self.accept(index)
2473+
if isinstance(typ, Instance) and typ.final_value is not None:
2474+
typ = typ.final_value
2475+
24712476
if isinstance(typ, LiteralType) and isinstance(typ.value, str):
24722477
item_name = typ.value
24732478
else:

mypy/plugins/common.py

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@
66
)
77
from mypy.plugin import ClassDefContext
88
from mypy.semanal import set_callable_name
9-
from mypy.types import CallableType, Overloaded, Type, TypeVarDef, LiteralType
9+
from mypy.types import CallableType, Overloaded, Type, TypeVarDef, LiteralType, Instance
1010
from mypy.typevars import fill_typevars
1111

1212

@@ -118,6 +118,9 @@ def try_getting_str_literal(expr: Expression, typ: Type) -> Optional[str]:
118118
"""If this expression is a string literal, or if the corresponding type
119119
is something like 'Literal["some string here"]', returns the underlying
120120
string value. Otherwise, returns None."""
121+
if isinstance(typ, Instance) and typ.final_value is not None:
122+
typ = typ.final_value
123+
121124
if isinstance(typ, LiteralType) and typ.fallback.type.fullname() == 'builtins.str':
122125
val = typ.value
123126
assert isinstance(val, str)

test-data/unit/check-literal.test

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2162,6 +2162,45 @@ del d[c_key] # E: TypedDict "Outer" has no key 'c'
21622162
[typing fixtures/typing-full.pyi]
21632163
[out]
21642164

2165+
[case testLiteralIntelligentIndexingUsingFinal]
2166+
from typing import Tuple, NamedTuple
2167+
from typing_extensions import Literal, Final
2168+
from mypy_extensions import TypedDict
2169+
2170+
int_key_good: Final = 0
2171+
int_key_bad: Final = 3
2172+
str_key_good: Final = "foo"
2173+
str_key_bad: Final = "missing"
2174+
2175+
class Unrelated: pass
2176+
2177+
MyTuple = NamedTuple('MyTuple', [
2178+
('foo', int),
2179+
('bar', str),
2180+
])
2181+
2182+
class MyDict(TypedDict):
2183+
foo: int
2184+
bar: str
2185+
2186+
a: Tuple[int, str]
2187+
b: MyTuple
2188+
c: MyDict
2189+
u: Unrelated
2190+
2191+
reveal_type(a[int_key_good]) # E: Revealed type is 'builtins.int'
2192+
reveal_type(b[int_key_good]) # E: Revealed type is 'builtins.int'
2193+
reveal_type(c[str_key_good]) # E: Revealed type is 'builtins.int'
2194+
reveal_type(c.get(str_key_good, u)) # E: Revealed type is 'Union[builtins.int, __main__.Unrelated]'
2195+
2196+
a[int_key_bad] # E: Tuple index out of range
2197+
b[int_key_bad] # E: Tuple index out of range
2198+
c[str_key_bad] # E: TypedDict "MyDict" has no key 'missing'
2199+
c.get(str_key_bad, u) # E: TypedDict "MyDict" has no key 'missing'
2200+
[builtins fixtures/dict.pyi]
2201+
[typing fixtures/typing-full.pyi]
2202+
[out]
2203+
21652204
[case testLiteralIntelligentIndexingTypedDictPython2-skip]
21662205
# flags: --python-version 2.7
21672206
from normal_mod import NormalDict

0 commit comments

Comments
 (0)