Disallow direct item access of NotRequired TypedDict properties: these should

always be accessed through .get() because the keys may not be present.

Fixes python#12094

Author: bsmedberg-xometry

mypy/checkexpr.py

@@ -2932,6 +2932,9 @@ def visit_unary_expr(self, e: UnaryExpr) -> Type:
     def visit_index_expr(self, e: IndexExpr) -> Type:
         """Type check an index expression (base[index]).
 
+        This function is only used for *expressions* (rvalues) not for setitem
+        statements (lvalues).
+
         It may also represent type application.
         """
         result = self.visit_index_expr_helper(e)
@@ -2988,7 +2991,7 @@ def visit_index_with_type(self, left_type: Type, e: IndexExpr,
             else:
                 return self.nonliteral_tuple_index_helper(left_type, index)
         elif isinstance(left_type, TypedDictType):
-            return self.visit_typeddict_index_expr(left_type, e.index)
+            return self.visit_typeddict_index_expr(left_type, e.index, is_expression=True)
         elif (isinstance(left_type, CallableType)
               and left_type.is_type_obj() and left_type.type_object().is_enum):
             return self.visit_enum_index_expr(left_type.type_object(), e.index, e)
@@ -3081,7 +3084,9 @@ def nonliteral_tuple_index_helper(self, left_type: TupleType, index: Expression)
 
     def visit_typeddict_index_expr(self, td_type: TypedDictType,
                                    index: Expression,
-                                   local_errors: Optional[MessageBuilder] = None
+                                   local_errors: Optional[MessageBuilder] = None,
+                                   *,
+                                   is_expression: bool
                                    ) -> Type:
         local_errors = local_errors or self.msg
         if isinstance(index, (StrExpr, UnicodeExpr)):
@@ -3113,6 +3118,8 @@ def visit_typeddict_index_expr(self, td_type: TypedDictType,
                 local_errors.typeddict_key_not_found(td_type, key_name, index)
                 return AnyType(TypeOfAny.from_error)
             else:
+                if is_expression and not td_type.is_required(key_name):
+                    local_errors.typeddict_key_not_required(td_type, key_name, index)
                 value_types.append(value_type)
         return make_simplified_union(value_types)
 

mypy/checkmember.py

@@ -856,7 +856,7 @@ def analyze_typeddict_access(name: str, typ: TypedDictType,
             # Since we can get this during `a['key'] = ...`
             # it is safe to assume that the context is `IndexExpr`.
             item_type = mx.chk.expr_checker.visit_typeddict_index_expr(
-                typ, mx.context.index)
+                typ, mx.context.index, is_expression=False)
         else:
             # It can also be `a.__setitem__(...)` direct call.
             # In this case `item_type` can be `Any`,

mypy/checkpattern.py

@@ -412,7 +412,7 @@ def get_mapping_item_type(self,
         mapping_type = get_proper_type(mapping_type)
         if isinstance(mapping_type, TypedDictType):
             result: Optional[Type] = self.chk.expr_checker.visit_typeddict_index_expr(
-                mapping_type, key, local_errors=local_errors)
+                mapping_type, key, local_errors=local_errors, is_expression=False)
             # If we can't determine the type statically fall back to treating it as a normal
             # mapping
             if local_errors.is_errors():

mypy/errorcodes.py

@@ -69,6 +69,9 @@ def __str__(self) -> str:
 TYPEDDICT_ITEM: Final = ErrorCode(
     "typeddict-item", "Check items when constructing TypedDict", "General"
 )
+TYPEDDICT_ITEM_ACCESS: Final = ErrorCode(
+    "typeddict-item-access", "Check item access when using TypedDict", "General"
+)
 HAS_TYPE: Final = ErrorCode(
     "has-type", "Check that type of reference can be determined", "General"
 )

mypy/messages.py

@@ -1276,6 +1276,17 @@ def typeddict_key_not_found(
                 self.note("Did you mean {}?".format(
                     pretty_seq(matches[:3], "or")), context, code=codes.TYPEDDICT_ITEM)
 
+    def typeddict_key_not_required(
+            self,
+            typ: TypedDictType,
+            item_name: str,
+            context: Context) -> None:
+        type_name: str = ""
+        if not typ.is_anonymous():
+            type_name = format_type(typ) + " "
+        self.fail('TypedDict {}key "{}" is not required.'.format(
+            type_name, item_name), context, code=codes.TYPEDDICT_ITEM_ACCESS)
+
     def typeddict_context_ambiguous(
             self,
             types: List[TypedDictType],

mypy/types.py

@@ -1680,6 +1680,9 @@ def __init__(self, items: 'OrderedDict[str, Type]', required_keys: Set[str],
     def accept(self, visitor: 'TypeVisitor[T]') -> T:
         return visitor.visit_typeddict_type(self)
 
+    def is_required(self, key: str) -> bool:
+        return key in self.required_keys
+
     def __hash__(self) -> int:
         return hash((frozenset(self.items.items()), self.fallback,
                      frozenset(self.required_keys)))

test-data/unit/check-typeddict.test

@@ -751,7 +751,6 @@ def get_coordinate(p: TaggedPoint, key: str) -> Union[str, int]:
     return p[key]  # E: TypedDict key must be a string literal; expected one of ("type", "x", "y")
 [builtins fixtures/dict.pyi]
 
-
 -- Special Method: __setitem__
 
 [case testCanSetItemOfTypedDictWithValidStringLiteralKeyAndCompatibleValueType]
@@ -1048,7 +1047,8 @@ reveal_type(d.get('x', {})) \
 reveal_type(d.get('x', None)) \
     # N: Revealed type is "Union[TypedDict('__main__.C', {'a': builtins.int}), None]"
 reveal_type(d.get('x', {}).get('a')) # N: Revealed type is "Union[builtins.int, None]"
-reveal_type(d.get('x', {})['a']) # N: Revealed type is "builtins.int"
+reveal_type(d.get('x', {})['a']) # N: Revealed type is "builtins.int" \
+                                 # E: TypedDict "C" key "a" is not required.
 [builtins fixtures/dict.pyi]
 [typing fixtures/typing-typeddict.pyi]
 
@@ -1100,8 +1100,10 @@ f(D(x='')) # E: Incompatible types (expression has type "str", TypedDict item "x
 from mypy_extensions import TypedDict
 D = TypedDict('D', {'x': int, 'y': str}, total=False)
 d: D
-reveal_type(d['x']) # N: Revealed type is "builtins.int"
-reveal_type(d['y']) # N: Revealed type is "builtins.str"
+reveal_type(d['x']) # N: Revealed type is "builtins.int" \
+                    # E: TypedDict "D" key "x" is not required.
+reveal_type(d['y']) # N: Revealed type is "builtins.str" \
+                    # E: TypedDict "D" key "y" is not required.
 reveal_type(d.get('x')) # N: Revealed type is "builtins.int"
 reveal_type(d.get('y')) # N: Revealed type is "builtins.str"
 [builtins fixtures/dict.pyi]
@@ -2284,6 +2286,36 @@ from typing import NotRequired
 Foo = TypedDict("Foo", {"a.x": NotRequired[int]})
 [typing fixtures/typing-typeddict.pyi]
 
+[case testCannotGetItemNotRequired]
+from typing import TypedDict
+from typing import NotRequired
+TaggedPoint = TypedDict('TaggedPoint', {'x': int, 'y': NotRequired[int]})
+p: TaggedPoint
+p['y']  # E: TypedDict "TaggedPoint" key "y" is not required.
+[typing fixtures/typing-typeddict.pyi]
+
+[case testCannotGetItemNotTotal]
+from typing import TypedDict
+TaggedPoint = TypedDict('TaggedPoint', {'x': int, 'y': int}, total=False)
+p: TaggedPoint
+p['y']  # E: TypedDict "TaggedPoint" key "y" is not required.
+[typing fixtures/typing-typeddict.pyi]
+
+[case testCanSetItemNotRequired]
+from typing import TypedDict
+from typing import NotRequired
+TaggedPoint = TypedDict('TaggedPoint', {'x': int, 'y': NotRequired[int]})
+p: TaggedPoint
+p['y'] = 1
+[typing fixtures/typing-typeddict.pyi]
+
+[case testCanSetItemNotTotal]
+from typing import TypedDict
+TaggedPoint = TypedDict('TaggedPoint', {'x': int, 'y': int}, total=False)
+p: TaggedPoint
+p['y'] = 1
+[typing fixtures/typing-typeddict.pyi]
+
 -- Union dunders
 
 [case testTypedDictUnionGetItem]