Merge pull request #69 from dawran6/issue-68

Fix the return type of escape to Markup

4 files changed, 23 insertions, 2 deletions

diff --git a/CHANGES b/CHANGES
index 706a230..ef0d801 100644
--- a/CHANGES
+++ b/CHANGES
@@ -1,6 +1,16 @@
 MarkupSafe Changelog
 ====================
 
+Version 1.1
+-----------
+
+unreleased
+
+- ``escape`` wraps ``__html__`` result in ``Markup``, consistent with
+  documented behavior. (`#69`_)
+
+.. _#69: https://github.com/pallets/markupsafe/pull/69
+
 Version 1.0
 -----------
 
diff --git a/markupsafe/_native.py b/markupsafe/_native.py
index 5e83f10..2c54a2d 100644
--- a/markupsafe/_native.py
+++ b/markupsafe/_native.py
@@ -18,7 +18,7 @@ def escape(s):
     such characters in HTML.  Marks return value as markup string.
     """
     if hasattr(s, '__html__'):
-        return s.__html__()
+        return Markup(s.__html__())
     return Markup(text_type(s)
         .replace('&', '&')
         .replace('>', '>')
diff --git a/markupsafe/_speedups.c b/markupsafe/_speedups.c
index d779a68..afe3425 100644
--- a/markupsafe/_speedups.c
+++ b/markupsafe/_speedups.c
@@ -131,8 +131,11 @@ escape(PyObject *self, PyObject *text)
 	/* if the object has an __html__ method that performs the escaping */
 	html = PyObject_GetAttrString(text, "__html__");
 	if (html) {
-		rv = PyObject_CallObject(html, NULL);
+		s = PyObject_CallObject(html, NULL);
 		Py_DECREF(html);
+		/* Convert to Markup object */
+		rv = PyObject_CallFunctionObjArgs(markup, (PyObject*)s, NULL);
+		Py_DECREF(s);
 		return rv;
 	}
 
diff --git a/tests.py b/tests.py
index 7ec2acb..ba5d30c 100755
--- a/tests.py
+++ b/tests.py
@@ -173,6 +173,14 @@ class MarkupTestCase(unittest.TestCase):
     def test_mul(self):
         self.assertEqual(Markup('a') * 3, Markup('aaa'))
 
+    def test_escape_return_type(self):
+        self.assertTrue(isinstance(escape('a'), Markup))
+        self.assertTrue(isinstance(escape(Markup('a')), Markup))
+        class Foo:
+            def __html__(self):
+                return '<strong>Foo</strong>'
+        self.assertTrue(isinstance(escape(Foo()), Markup))
+
 
 class MarkupLeakTestCase(unittest.TestCase):