From 791e598e3dceca805e6fa40e2dcfffd300963780 Mon Sep 17 00:00:00 2001
From: Jeong YunWon <jeong@youknowone.org>
Date: Fri, 11 Oct 2019 00:21:19 +0900
Subject: [PATCH]  Fix builtin_round with non-int __round__

---
 tests/snippets/builtin_round.py | 10 ++++++++++
 vm/src/builtins.rs              | 16 ++++++----------
 2 files changed, 16 insertions(+), 10 deletions(-)

diff --git a/tests/snippets/builtin_round.py b/tests/snippets/builtin_round.py
index 5526c551e..bad9c9e05 100644
--- a/tests/snippets/builtin_round.py
+++ b/tests/snippets/builtin_round.py
@@ -19,3 +19,13 @@ with assert_raises(TypeError):
     round(0, 0.0)
 with assert_raises(TypeError):
     round(0.0, 0.0)
+
+
+class X:
+    def __round__(self, ndigits=None):
+        return 1.1
+
+
+assert round(X(), 1) == 1.1
+assert round(X(), None) == 1.1
+assert round(X()) == 1.1
diff --git a/vm/src/builtins.rs b/vm/src/builtins.rs
index a1153e597..1584b5aa5 100644
--- a/vm/src/builtins.rs
+++ b/vm/src/builtins.rs
@@ -674,24 +674,20 @@ fn builtin_round(
     ndigits: OptionalArg<Option<PyIntRef>>,
     vm: &VirtualMachine,
 ) -> PyResult {
-    match ndigits {
+    let rounded = match ndigits {
         OptionalArg::Present(ndigits) => match ndigits {
             Some(int) => {
                 let ndigits = vm.call_method(int.as_object(), "__int__", vec![])?;
-                let rounded = vm.call_method(&number, "__round__", vec![ndigits])?;
-                Ok(rounded)
-            }
-            None => {
-                let rounded = &vm.call_method(&number, "__round__", vec![])?;
-                Ok(vm.ctx.new_int(objint::get_value(rounded).clone()))
+                vm.call_method(&number, "__round__", vec![ndigits])?
             }
+            None => vm.call_method(&number, "__round__", vec![])?,
         },
         OptionalArg::Missing => {
             // without a parameter, the result type is coerced to int
-            let rounded = &vm.call_method(&number, "__round__", vec![])?;
-            Ok(vm.ctx.new_int(objint::get_value(rounded).clone()))
+            vm.call_method(&number, "__round__", vec![])?
         }
-    }
+    };
+    Ok(rounded)
 }
 
 fn builtin_setattr(