Merge branch 'master' into dict-literals

Author: BobTheBuidler

docs/source/generics.rst

@@ -1165,7 +1165,7 @@ This example correctly uses a covariant type variable:
 
 See :ref:`variance-of-generics` for more about variance.
 
-Generic protocols can also be recursive. Example (Python 3.12 synta):
+Generic protocols can also be recursive. Example (Python 3.12 syntax):
 
 .. code-block:: python
 

mypy/checkexpr.py

@@ -1502,7 +1502,7 @@ def check_call_expr_with_callee_type(
     def check_union_call_expr(self, e: CallExpr, object_type: UnionType, member: str) -> Type:
         """Type check calling a member expression where the base type is a union."""
         res: list[Type] = []
-        for typ in object_type.relevant_items():
+        for typ in flatten_nested_unions(object_type.relevant_items()):
             # Member access errors are already reported when visiting the member expression.
             with self.msg.filter_errors():
                 item = analyze_member_access(

mypy/errors.py

@@ -1003,6 +1003,9 @@ def format_messages(
                     marker = "^"
                     if end_line == line and end_column > column:
                         marker = f'^{"~" * (end_column - column - 1)}'
+                    elif end_line != line:
+                        # just highlight the first line instead
+                        marker = f'^{"~" * (len(source_line_expanded) - column - 1)}'
                     a.append(" " * (DEFAULT_SOURCE_OFFSET + column) + marker)
         return a
 

mypy/plugins/default.py

@@ -5,7 +5,7 @@
 
 import mypy.errorcodes as codes
 from mypy import message_registry
-from mypy.nodes import DictExpr, IntExpr, StrExpr, UnaryExpr
+from mypy.nodes import DictExpr, Expression, IntExpr, StrExpr, UnaryExpr
 from mypy.plugin import (
     AttributeContext,
     ClassDefContext,
@@ -263,30 +263,40 @@ def typed_dict_get_callback(ctx: MethodContext) -> Type:
         if keys is None:
             return ctx.default_return_type
 
+        default_type: Type
+        default_arg: Expression | None
+        if len(ctx.arg_types) <= 1 or not ctx.arg_types[1]:
+            default_arg = None
+            default_type = NoneType()
+        elif len(ctx.arg_types[1]) == 1 and len(ctx.args[1]) == 1:
+            default_arg = ctx.args[1][0]
+            default_type = ctx.arg_types[1][0]
+        else:
+            return ctx.default_return_type
+
         output_types: list[Type] = []
         for key in keys:
-            value_type = get_proper_type(ctx.type.items.get(key))
+            value_type: Type | None = ctx.type.items.get(key)
             if value_type is None:
                 return ctx.default_return_type
 
-            if len(ctx.arg_types) == 1:
+            if key in ctx.type.required_keys:
                 output_types.append(value_type)
-            elif len(ctx.arg_types) == 2 and len(ctx.arg_types[1]) == 1 and len(ctx.args[1]) == 1:
-                default_arg = ctx.args[1][0]
+            else:
+                # HACK to deal with get(key, {})
                 if (
                     isinstance(default_arg, DictExpr)
                     and len(default_arg.items) == 0
-                    and isinstance(value_type, TypedDictType)
+                    and isinstance(vt := get_proper_type(value_type), TypedDictType)
                 ):
-                    # Special case '{}' as the default for a typed dict type.
-                    output_types.append(value_type.copy_modified(required_keys=set()))
+                    output_types.append(vt.copy_modified(required_keys=set()))
                 else:
                     output_types.append(value_type)
-                    output_types.append(ctx.arg_types[1][0])
-
-        if len(ctx.arg_types) == 1:
-            output_types.append(NoneType())
+                    output_types.append(default_type)
 
+        # for nicer reveal_type, put default at the end, if it is present
+        if default_type in output_types:
+            output_types = [t for t in output_types if t != default_type] + [default_type]
         return make_simplified_union(output_types)
     return ctx.default_return_type
 

mypy/stubtest.py

@@ -968,7 +968,7 @@ def from_overloadedfuncdef(stub: nodes.OverloadedFuncDef) -> Signature[nodes.Arg
 
         is_arg_pos_only: defaultdict[str, set[bool]] = defaultdict(set)
         for func in map(_resolve_funcitem_from_decorator, stub.items):
-            assert func is not None, "Failed to resolve decorated overload"
+            assert func is not None, f"Failed to resolve decorated overload of {stub.fullname!r}"
             args = maybe_strip_cls(stub.name, func.arguments)
             for index, arg in enumerate(args):
                 if (
@@ -984,7 +984,7 @@ def from_overloadedfuncdef(stub: nodes.OverloadedFuncDef) -> Signature[nodes.Arg
 
         all_args: dict[str, list[tuple[nodes.Argument, int]]] = {}
         for func in map(_resolve_funcitem_from_decorator, stub.items):
-            assert func is not None, "Failed to resolve decorated overload"
+            assert func is not None, f"Failed to resolve decorated overload of {stub.fullname!r}"
             args = maybe_strip_cls(stub.name, func.arguments)
             for index, arg in enumerate(args):
                 # For positional-only args, we allow overloads to have different names for the same

mypyc/codegen/emitclass.py

@@ -410,7 +410,7 @@ def setter_name(cl: ClassIR, attribute: str, names: NameGenerator) -> str:
 
 
 def generate_object_struct(cl: ClassIR, emitter: Emitter) -> None:
-    seen_attrs: set[tuple[str, RType]] = set()
+    seen_attrs: set[str] = set()
     lines: list[str] = []
     lines += ["typedef struct {", "PyObject_HEAD", "CPyVTableItem *vtable;"]
     if cl.has_method("__call__"):
@@ -427,9 +427,11 @@ def generate_object_struct(cl: ClassIR, emitter: Emitter) -> None:
                             lines.append(f"{BITMAP_TYPE} {attr};")
                             bitmap_attrs.append(attr)
             for attr, rtype in base.attributes.items():
-                if (attr, rtype) not in seen_attrs:
+                # Generated class may redefine certain attributes with different
+                # types in subclasses (this would be unsafe for user-defined classes).
+                if attr not in seen_attrs:
                     lines.append(f"{emitter.ctype_spaced(rtype)}{emitter.attr(attr)};")
-                    seen_attrs.add((attr, rtype))
+                    seen_attrs.add(attr)
 
                     if isinstance(rtype, RTuple):
                         emitter.declare_tuple_struct(rtype)

mypyc/codegen/emitmodule.py

@@ -1067,6 +1067,8 @@ def emit_module_exec_func(
                     "(PyObject *){t}_template, NULL, modname);".format(t=type_struct)
                 )
                 emitter.emit_lines(f"if (unlikely(!{type_struct}))", "    goto fail;")
+                name_prefix = cl.name_prefix(emitter.names)
+                emitter.emit_line(f"CPyDef_{name_prefix}_trait_vtable_setup();")
 
         emitter.emit_lines("if (CPyGlobalsInit() < 0)", "    goto fail;")
 

mypyc/ir/func_ir.py

@@ -149,8 +149,11 @@ def __init__(
         module_name: str,
         sig: FuncSignature,
         kind: int = FUNC_NORMAL,
+        *,
         is_prop_setter: bool = False,
         is_prop_getter: bool = False,
+        is_generator: bool = False,
+        is_coroutine: bool = False,
         implicit: bool = False,
         internal: bool = False,
     ) -> None:
@@ -161,6 +164,8 @@ def __init__(
         self.kind = kind
         self.is_prop_setter = is_prop_setter
         self.is_prop_getter = is_prop_getter
+        self.is_generator = is_generator
+        self.is_coroutine = is_coroutine
         if class_name is None:
             self.bound_sig: FuncSignature | None = None
         else:
@@ -219,6 +224,8 @@ def serialize(self) -> JsonDict:
             "kind": self.kind,
             "is_prop_setter": self.is_prop_setter,
             "is_prop_getter": self.is_prop_getter,
+            "is_generator": self.is_generator,
+            "is_coroutine": self.is_coroutine,
             "implicit": self.implicit,
             "internal": self.internal,
         }
@@ -240,10 +247,12 @@ def deserialize(cls, data: JsonDict, ctx: DeserMaps) -> FuncDecl:
             data["module_name"],
             FuncSignature.deserialize(data["sig"], ctx),
             data["kind"],
-            data["is_prop_setter"],
-            data["is_prop_getter"],
-            data["implicit"],
-            data["internal"],
+            is_prop_setter=data["is_prop_setter"],
+            is_prop_getter=data["is_prop_getter"],
+            is_generator=data["is_generator"],
+            is_coroutine=data["is_coroutine"],
+            implicit=data["implicit"],
+            internal=data["internal"],
         )
 
 

mypyc/irbuild/context.py

@@ -98,7 +98,11 @@ def curr_env_reg(self) -> Value:
     def can_merge_generator_and_env_classes(self) -> bool:
         # In simple cases we can place the environment into the generator class,
         # instead of having two separate classes.
-        return self.is_generator and not self.is_nested and not self.contains_nested
+        if self._generator_class and not self._generator_class.ir.is_final_class:
+            result = False
+        else:
+            result = self.is_generator and not self.is_nested and not self.contains_nested
+        return result
 
 
 class ImplicitClass:

mypyc/irbuild/for_helpers.py

@@ -1203,18 +1203,18 @@ def gen_cleanup(self) -> None:
             gen.gen_cleanup()
 
 
-def get_expr_length(expr: Expression) -> int | None:
+def get_expr_length(builder: IRBuilder, expr: Expression) -> int | None:
     if isinstance(expr, (StrExpr, BytesExpr)):
         return len(expr.value)
     elif isinstance(expr, (ListExpr, TupleExpr)):
         # if there are no star expressions, or we know the length of them,
         # we know the length of the expression
-        stars = [get_expr_length(i) for i in expr.items if isinstance(i, StarExpr)]
+        stars = [get_expr_length(builder, i) for i in expr.items if isinstance(i, StarExpr)]
         if None not in stars:
             other = sum(not isinstance(i, StarExpr) for i in expr.items)
             return other + sum(stars)  # type: ignore [arg-type]
     elif isinstance(expr, StarExpr):
-        return get_expr_length(expr.expr)
+        return get_expr_length(builder, expr.expr)
     elif (
         isinstance(expr, RefExpr)
         and isinstance(expr.node, Var)
@@ -1227,6 +1227,11 @@ def get_expr_length(expr: Expression) -> int | None:
     # performance boost and can be (sometimes) figured out pretty easily. set and dict
     # comps *can* be done as well but will need special logic to consider the possibility
     # of key conflicts. Range, enumerate, zip are all simple logic.
+
+    # we might still be able to get the length directly from the type
+    rtype = builder.node_type(expr)
+    if isinstance(rtype, RTuple):
+        return len(rtype.types)
     return None
 
 
@@ -1235,7 +1240,7 @@ def get_expr_length_value(
 ) -> Value:
     rtype = builder.node_type(expr)
     assert is_sequence_rprimitive(rtype) or isinstance(rtype, RTuple), rtype
-    length = get_expr_length(expr)
+    length = get_expr_length(builder, expr)
     if length is None:
         # We cannot compute the length at compile time, so we will fetch it.
         return builder.builder.builtin_len(expr_reg, line, use_pyssize_t=use_pyssize_t)