Also support optional bytes values

Author: JukkaL

mypyc/irbuild/ll_builder.py

@@ -2522,20 +2522,19 @@ def translate_eq_cmp(self, lreg: Value, rreg: Value, expr_op: str, line: int) ->
 
         lopt = optional_value_type(ltype)
         ropt = optional_value_type(rtype)
+
+        # Can we do a quick comparison of two optional types (special case None values)?
         fast_opt_eq = False
-        if lopt is not None and ropt is not None:
-            # Can we do a quick comparison of two optional types?
-            if is_same_type(lopt, ropt) and is_str_rprimitive(lopt):
+        if lopt is not None:
+            if ropt is not None and is_same_type(lopt, ropt) and self._never_equal_to_none(lopt):
                 fast_opt_eq = True
-        elif lopt is not None:
-            if is_same_type(lopt, rtype) and is_str_rprimitive(lopt):
+            if is_same_type(lopt, rtype) and self._never_equal_to_none(lopt):
                 fast_opt_eq = True
         elif ropt is not None:
-            if is_same_type(ropt, ltype) and is_str_rprimitive(ropt):
+            if is_same_type(ropt, ltype) and self._never_equal_to_none(ropt):
                 fast_opt_eq = True
-
         if fast_opt_eq:
-            return self.translate_fast_optional_eq_cmp(lreg, rreg, expr_op, line)
+            return self._translate_fast_optional_eq_cmp(lreg, rreg, expr_op, line)
 
         if not (isinstance(ltype, RInstance) and ltype == rtype):
             return None
@@ -2563,11 +2562,18 @@ def translate_eq_cmp(self, lreg: Value, rreg: Value, expr_op: str, line: int) ->
 
         return self.gen_method_call(lreg, op_methods[expr_op], [rreg], ltype, line)
 
-    def translate_fast_optional_eq_cmp(
+    def _never_equal_to_none(self, typ: RType) -> bool:
+        """Are the values of type never equal to None?"""
+        # TODO: Support RInstance with no custom __eq__/__ne__ and other primitive types.
+        return is_str_rprimitive(typ) or is_bytes_rprimitive(typ)
+
+    def _translate_fast_optional_eq_cmp(
         self, lreg: Value, rreg: Value, expr_op: str, line: int
     ) -> Value:
         if not isinstance(lreg.type, RUnion):
             lreg, rreg = rreg, lreg
+        value_typ = optional_value_type(lreg.type)
+        assert value_typ
         res = Register(bool_rprimitive)
         x = self.add(ComparisonOp(lreg, self.none_object(), ComparisonOp.EQ, line))
         l_none = BasicBlock()
@@ -2584,12 +2590,13 @@ def translate_fast_optional_eq_cmp(
         self.goto(out)
         self.activate_block(l_not_none)
         if not isinstance(rreg.type, RUnion):
-            z = self.compare_strings(
-                self.unbox_or_cast(lreg, str_rprimitive, line, can_borrow=True, unchecked=True),
+            z = self.translate_eq_cmp(
+                self.unbox_or_cast(lreg, value_typ, line, can_borrow=True, unchecked=True),
                 rreg,
                 expr_op,
                 line,
             )
+            assert z is not None
             self.add(Assign(res, z))
         else:
             r_none = BasicBlock()
@@ -2600,12 +2607,13 @@ def translate_fast_optional_eq_cmp(
             self.add(Assign(res, self.false()))
             self.goto(out)
             self.activate_block(r_not_none)
-            z = self.compare_strings(
-                self.unbox_or_cast(lreg, str_rprimitive, line, can_borrow=True, unchecked=True),
-                self.unbox_or_cast(rreg, str_rprimitive, line, can_borrow=True, unchecked=True),
+            z = self.translate_eq_cmp(
+                self.unbox_or_cast(lreg, value_typ, line, can_borrow=True, unchecked=True),
+                self.unbox_or_cast(rreg, value_typ, line, can_borrow=True, unchecked=True),
                 expr_op,
                 line,
             )
+            assert z is not None
             self.add(Assign(res, z))
         self.goto(out)
         self.activate_block(out)

mypyc/test-data/irbuild-str.test

@@ -741,20 +741,21 @@ L3:
     keep_alive x
     return r2
 
-[case testOptionalStrEquality3]
+[case testOptionalBytesEquality]
 from typing import Optional
 
-def non_opt_opt(x: str, y: Optional[str]) -> bool:
+def non_opt_opt(x: bytes, y: Optional[bytes]) -> bool:
     return x == y
 [out]
 def non_opt_opt(x, y):
-    x :: str
-    y :: union[str, None]
+    x :: bytes
+    y :: union[bytes, None]
     r0 :: object
     r1 :: bit
     r2 :: bool
-    r3 :: str
-    r4 :: bool
+    r3 :: bytes
+    r4 :: i32
+    r5, r6 :: bit
 L0:
     r0 = load_address _Py_NoneStruct
     r1 = y == r0
@@ -763,9 +764,11 @@ L1:
     r2 = 0
     goto L3
 L2:
-    r3 = unchecked borrow cast(str, y)
-    r4 = CPyStr_Equal(r3, x)
-    r2 = r4
+    r3 = unchecked borrow cast(bytes, y)
+    r4 = CPyBytes_Compare(r3, x)
+    r5 = r4 >= 0 :: signed
+    r6 = r4 == 1
+    r2 = r6
 L3:
     keep_alive y
     return r2