checkexpr: broader fast path for literal container expressions

The original implementation was very conservative and bailed whenever a type context
was present, drastically limiting its reach.

There are a number of type contexts where the fast path can safely be taken.
Most importantly:
 - Any
 - matching container types

Author: hugues-aff

mypy/checkexpr.py

@@ -98,7 +98,13 @@
 )
 from mypy.semanal_enum import ENUM_BASES
 from mypy.state import state
-from mypy.subtypes import is_equivalent, is_same_type, is_subtype, non_method_protocol_members
+from mypy.subtypes import (
+    is_equivalent,
+    is_proper_subtype,
+    is_same_type,
+    is_subtype,
+    non_method_protocol_members,
+)
 from mypy.traverser import has_await_expression
 from mypy.typeanal import (
     check_for_explicit_any,
@@ -3994,18 +4000,36 @@ def fast_container_type(
         self, e: ListExpr | SetExpr | TupleExpr, container_fullname: str
     ) -> Type | None:
         """
-        Fast path to determine the type of a list or set literal,
-        based on the list of entries. This mostly impacts large
-        module-level constant definitions.
+        Fast path to determine the type of a list or set literal, based on the list of entries.
+        This mostly impacts large constant definitions.
 
         Limitations:
          - no active type context
          - no star expressions
          - the joined type of all entries must be an Instance or Tuple type
         """
-        ctx = self.type_context[-1]
-        if ctx:
+        ctx = get_proper_type(self.type_context[-1])
+        # TODO: can we safely allow TypeVarType (with appropriate upper_bound or values) ?
+        if not ctx or isinstance(ctx, AnyType):
+            return self.fast_container_type_no_context(e, container_fullname)
+
+        # TODO: support relevant typing base classes (Sequence, AbstractSet, ...) ?
+        if not (isinstance(ctx, Instance) and ctx.type.fullname == container_fullname):
             return None
+
+        vt = get_proper_type(ctx.args[0])
+        if not (isinstance(vt, Instance) and allow_fast_container_literal(vt)):
+            return None
+
+        rt = self.fast_container_type_no_context(e, container_fullname)
+        if rt and is_proper_subtype(rt.args[0], vt):
+            # let the type context win if we inferred a more precise type (List is invariant...)
+            return ctx
+        return None
+
+    def fast_container_type_no_context(
+        self, e: ListExpr | SetExpr | TupleExpr, container_fullname: str
+    ) -> Instance | None:
         rt = self.resolved_type.get(e, None)
         if rt is not None:
             return rt if isinstance(rt, Instance) else None
@@ -4116,18 +4140,48 @@ def visit_tuple_expr(self, e: TupleExpr) -> Type:
 
     def fast_dict_type(self, e: DictExpr) -> Type | None:
         """
-        Fast path to determine the type of a dict literal,
-        based on the list of entries. This mostly impacts large
-        module-level constant definitions.
+        Fast path to determine the type of a dict literal, based on the list of entries.
+        This mostly impacts large module-level constant definitions.
+
+        Will only trigger for the following type contexts:
+         - None (i.e no type context)
+         - AnyType
+         - Dict[K, V] where K and V are both Instance or Tuple types
+        """
+        ctx = get_proper_type(self.type_context[-1])
+        # TODO: can we safely allow TypeVarType (with appropriate upper_bound or values) ?
+        if not ctx or isinstance(ctx, AnyType):
+            return self.fast_dict_type_no_context(e)
+
+        # TODO: support relevant typing base classes (Mapping, MutableMapping, ...) ?
+        if not (isinstance(ctx, Instance) and ctx.type.fullname == "builtins.dict"):
+            return None
+
+        kt = get_proper_type(ctx.args[0])
+        vt = get_proper_type(ctx.args[1])
+        if not (
+            isinstance(kt, Instance)
+            and isinstance(vt, Instance)
+            and allow_fast_container_literal(kt)
+            and allow_fast_container_literal(vt)
+        ):
+            return None
+
+        rt = self.fast_dict_type_no_context(e)
+        if rt and is_proper_subtype(rt.args[0], kt) and is_proper_subtype(rt.args[1], vt):
+            # let the type context win if we inferred a more precise type (Dict is invariant...)
+            return ctx
+        return None
+
+    def fast_dict_type_no_context(self, e: DictExpr) -> Instance | None:
+        """
+        Fast path to determine the type of a dict literal, based on the list of entries.
+        This mostly impacts large constant definitions.
 
         Limitations:
-         - no active type context
          - only supported star expressions are other dict instances
          - the joined types of all keys and values must be Instance or Tuple types
         """
-        ctx = self.type_context[-1]
-        if ctx:
-            return None
         rt = self.resolved_type.get(e, None)
         if rt is not None:
             return rt if isinstance(rt, Instance) else None