diff --git a/mypy/checkexpr.py b/mypy/checkexpr.py index b6bb5f95662d..9d29c1bd0bba 100644 --- a/mypy/checkexpr.py +++ b/mypy/checkexpr.py @@ -2459,6 +2459,8 @@ def _get_value(self, index: Expression) -> Optional[int]: if isinstance(operand, IntExpr): return -1 * operand.value typ = self.accept(index) + if isinstance(typ, Instance) and typ.final_value is not None: + typ = typ.final_value if isinstance(typ, LiteralType) and isinstance(typ.value, int): return typ.value return None @@ -2468,6 +2470,9 @@ def visit_typeddict_index_expr(self, td_type: TypedDictType, index: Expression) item_name = index.value else: typ = self.accept(index) + if isinstance(typ, Instance) and typ.final_value is not None: + typ = typ.final_value + if isinstance(typ, LiteralType) and isinstance(typ.value, str): item_name = typ.value else: diff --git a/mypy/plugins/common.py b/mypy/plugins/common.py index c1dcd6b4ca2e..e95634152ab0 100644 --- a/mypy/plugins/common.py +++ b/mypy/plugins/common.py @@ -6,7 +6,7 @@ ) from mypy.plugin import ClassDefContext from mypy.semanal import set_callable_name -from mypy.types import CallableType, Overloaded, Type, TypeVarDef, LiteralType +from mypy.types import CallableType, Overloaded, Type, TypeVarDef, LiteralType, Instance from mypy.typevars import fill_typevars @@ -118,6 +118,9 @@ def try_getting_str_literal(expr: Expression, typ: Type) -> Optional[str]: """If this expression is a string literal, or if the corresponding type is something like 'Literal["some string here"]', returns the underlying string value. Otherwise, returns None.""" + if isinstance(typ, Instance) and typ.final_value is not None: + typ = typ.final_value + if isinstance(typ, LiteralType) and typ.fallback.type.fullname() == 'builtins.str': val = typ.value assert isinstance(val, str) diff --git a/test-data/unit/check-literal.test b/test-data/unit/check-literal.test index a57f93f7d9ab..ab94333ce81a 100644 --- a/test-data/unit/check-literal.test +++ b/test-data/unit/check-literal.test @@ -2162,6 +2162,45 @@ del d[c_key] # E: TypedDict "Outer" has no key 'c' [typing fixtures/typing-full.pyi] [out] +[case testLiteralIntelligentIndexingUsingFinal] +from typing import Tuple, NamedTuple +from typing_extensions import Literal, Final +from mypy_extensions import TypedDict + +int_key_good: Final = 0 +int_key_bad: Final = 3 +str_key_good: Final = "foo" +str_key_bad: Final = "missing" + +class Unrelated: pass + +MyTuple = NamedTuple('MyTuple', [ + ('foo', int), + ('bar', str), +]) + +class MyDict(TypedDict): + foo: int + bar: str + +a: Tuple[int, str] +b: MyTuple +c: MyDict +u: Unrelated + +reveal_type(a[int_key_good]) # E: Revealed type is 'builtins.int' +reveal_type(b[int_key_good]) # E: Revealed type is 'builtins.int' +reveal_type(c[str_key_good]) # E: Revealed type is 'builtins.int' +reveal_type(c.get(str_key_good, u)) # E: Revealed type is 'Union[builtins.int, __main__.Unrelated]' + +a[int_key_bad] # E: Tuple index out of range +b[int_key_bad] # E: Tuple index out of range +c[str_key_bad] # E: TypedDict "MyDict" has no key 'missing' +c.get(str_key_bad, u) # E: TypedDict "MyDict" has no key 'missing' +[builtins fixtures/dict.pyi] +[typing fixtures/typing-full.pyi] +[out] + [case testLiteralIntelligentIndexingTypedDictPython2-skip] # flags: --python-version 2.7 from normal_mod import NormalDict