Merge branch 'master' into constant-fold-namedtuple-fields

Author: BobTheBuidler

mypy/cache.py

@@ -1,65 +1,22 @@
 from __future__ import annotations
 
 from collections.abc import Sequence
-from typing import TYPE_CHECKING, Final
+from typing import Final
 
 from mypy_extensions import u8
-
-try:
-    from native_internal import (
-        Buffer as Buffer,
-        read_bool as read_bool,
-        read_float as read_float,
-        read_int as read_int,
-        read_str as read_str,
-        read_tag as read_tag,
-        write_bool as write_bool,
-        write_float as write_float,
-        write_int as write_int,
-        write_str as write_str,
-        write_tag as write_tag,
-    )
-except ImportError:
-    # TODO: temporary, remove this after we publish mypy-native on PyPI.
-    if not TYPE_CHECKING:
-
-        class Buffer:
-            def __init__(self, source: bytes = b"") -> None:
-                raise NotImplementedError
-
-            def getvalue(self) -> bytes:
-                raise NotImplementedError
-
-        def read_int(data: Buffer) -> int:
-            raise NotImplementedError
-
-        def write_int(data: Buffer, value: int) -> None:
-            raise NotImplementedError
-
-        def read_tag(data: Buffer) -> u8:
-            raise NotImplementedError
-
-        def write_tag(data: Buffer, value: u8) -> None:
-            raise NotImplementedError
-
-        def read_str(data: Buffer) -> str:
-            raise NotImplementedError
-
-        def write_str(data: Buffer, value: str) -> None:
-            raise NotImplementedError
-
-        def read_bool(data: Buffer) -> bool:
-            raise NotImplementedError
-
-        def write_bool(data: Buffer, value: bool) -> None:
-            raise NotImplementedError
-
-        def read_float(data: Buffer) -> float:
-            raise NotImplementedError
-
-        def write_float(data: Buffer, value: float) -> None:
-            raise NotImplementedError
-
+from native_internal import (
+    Buffer as Buffer,
+    read_bool as read_bool,
+    read_float as read_float,
+    read_int as read_int,
+    read_str as read_str,
+    read_tag as read_tag,
+    write_bool as write_bool,
+    write_float as write_float,
+    write_int as write_int,
+    write_str as write_str,
+    write_tag as write_tag,
+)
 
 # Always use this type alias to refer to type tags.
 Tag = u8

mypy/checkexpr.py

@@ -18,6 +18,7 @@
 from mypy.checker_shared import ExpressionCheckerSharedApi
 from mypy.checkmember import analyze_member_access, has_operator
 from mypy.checkstrformat import StringFormatterChecker
+from mypy.constant_fold import constant_fold_expr
 from mypy.erasetype import erase_type, remove_instance_last_known_values, replace_meta_vars
 from mypy.errors import ErrorInfo, ErrorWatcher, report_internal_error
 from mypy.expandtype import (
@@ -656,11 +657,12 @@ def visit_call_expr_inner(self, e: CallExpr, allow_none_return: bool = False) ->
         return ret_type
 
     def check_str_format_call(self, e: CallExpr) -> None:
-        """More precise type checking for str.format() calls on literals."""
+        """More precise type checking for str.format() calls on literals and folded constants."""
         assert isinstance(e.callee, MemberExpr)
         format_value = None
-        if isinstance(e.callee.expr, StrExpr):
-            format_value = e.callee.expr.value
+        folded_callee_expr = constant_fold_expr(e.callee.expr, "<unused>")
+        if isinstance(folded_callee_expr, str):
+            format_value = folded_callee_expr
         elif self.chk.has_type(e.callee.expr):
             typ = get_proper_type(self.chk.lookup_type(e.callee.expr))
             if (

mypyc/irbuild/expression.py

@@ -558,7 +558,7 @@ def transform_op_expr(builder: IRBuilder, expr: OpExpr) -> Value:
     # Special case some int ops to allow borrowing operands.
     if is_int_rprimitive(ltype) and is_int_rprimitive(rtype):
         if expr.op == "//":
-            expr = try_optimize_int_floor_divide(expr)
+            expr = try_optimize_int_floor_divide(builder, expr)
         if expr.op in int_borrow_friendly_op:
             borrow_left = is_borrow_friendly_expr(builder, expr.right)
             borrow_right = True
@@ -571,11 +571,11 @@ def transform_op_expr(builder: IRBuilder, expr: OpExpr) -> Value:
     return builder.binary_op(left, right, expr.op, expr.line)
 
 
-def try_optimize_int_floor_divide(expr: OpExpr) -> OpExpr:
+def try_optimize_int_floor_divide(builder: IRBuilder, expr: OpExpr) -> OpExpr:
     """Replace // with a power of two with a right shift, if possible."""
-    if not isinstance(expr.right, IntExpr):
+    divisor = constant_fold_expr(builder, expr.right)
+    if not isinstance(divisor, int):
         return expr
-    divisor = expr.right.value
     shift = divisor.bit_length() - 1
     if 0 < shift < 28 and divisor == (1 << shift):
         return OpExpr(">>", expr.left, IntExpr(shift))

mypyc/irbuild/specialize.py

@@ -78,6 +78,7 @@
     uint8_rprimitive,
 )
 from mypyc.irbuild.builder import IRBuilder
+from mypyc.irbuild.constant_fold import constant_fold_expr
 from mypyc.irbuild.for_helpers import (
     comprehension_helper,
     sequence_from_generator_preallocate_helper,
@@ -716,21 +717,18 @@ def translate_dict_setdefault(builder: IRBuilder, expr: CallExpr, callee: RefExp
 
 @specialize_function("format", str_rprimitive)
 def translate_str_format(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Value | None:
-    if (
-        isinstance(callee, MemberExpr)
-        and isinstance(callee.expr, StrExpr)
-        and expr.arg_kinds.count(ARG_POS) == len(expr.arg_kinds)
-    ):
-        format_str = callee.expr.value
-        tokens = tokenizer_format_call(format_str)
-        if tokens is None:
-            return None
-        literals, format_ops = tokens
-        # Convert variables to strings
-        substitutions = convert_format_expr_to_str(builder, format_ops, expr.args, expr.line)
-        if substitutions is None:
-            return None
-        return join_formatted_strings(builder, literals, substitutions, expr.line)
+    if isinstance(callee, MemberExpr):
+        folded_callee = constant_fold_expr(builder, callee.expr)
+        if isinstance(folded_callee, str) and expr.arg_kinds.count(ARG_POS) == len(expr.arg_kinds):
+            tokens = tokenizer_format_call(folded_callee)
+            if tokens is None:
+                return None
+            literals, format_ops = tokens
+            # Convert variables to strings
+            substitutions = convert_format_expr_to_str(builder, format_ops, expr.args, expr.line)
+            if substitutions is None:
+                return None
+            return join_formatted_strings(builder, literals, substitutions, expr.line)
     return None