py: Detect ZeroDivisionError properly for floats.

1 files changed, 9 insertions, 6 deletions

diff --git a/py/objfloat.c b/py/objfloat.c
index 3f4a68093..16d4fbfd5 100644
--- a/py/objfloat.c
+++ b/py/objfloat.c
@@ -118,15 +118,18 @@ mp_obj_t mp_obj_float_binary_op(int op, mp_float_t lhs_val, mp_obj_t rhs_in) {
         // TODO: verify that C floor matches Python semantics
         case MP_BINARY_OP_FLOOR_DIVIDE:
         case MP_BINARY_OP_INPLACE_FLOOR_DIVIDE:
+            if (rhs_val == 0) {
+                zero_division_error:
+                nlr_raise(mp_obj_new_exception_msg(&mp_type_ZeroDivisionError, "float division by zero"));
+            }
             lhs_val = MICROPY_FLOAT_C_FUN(floor)(lhs_val / rhs_val);
-            goto check_zero_division;
+            break;
         case MP_BINARY_OP_TRUE_DIVIDE:
-        case MP_BINARY_OP_INPLACE_TRUE_DIVIDE: 
-            lhs_val /= rhs_val; 
-check_zero_division:
-            if (isinf(lhs_val)){ // check for division by zero
-                nlr_raise(mp_obj_new_exception_msg(&mp_type_ZeroDivisionError, "float division by zero"));
+        case MP_BINARY_OP_INPLACE_TRUE_DIVIDE:
+            if (rhs_val == 0) {
+                goto zero_division_error;
             }
+            lhs_val /= rhs_val;
             break;
         case MP_BINARY_OP_POWER:
         case MP_BINARY_OP_INPLACE_POWER: lhs_val = MICROPY_FLOAT_C_FUN(pow)(lhs_val, rhs_val); break;