bpo-32363: Disable Task.set_exception() and Task.set_result() (#4923)

4 files changed, 94 insertions, 33 deletions

diff --git a/Lib/asyncio/futures.py b/Lib/asyncio/futures.py
index 24843c0..13fb47c 100644
--- a/Lib/asyncio/futures.py
+++ b/Lib/asyncio/futures.py
@@ -239,14 +239,15 @@ class Future:
         self._schedule_callbacks()
         self._log_traceback = True
 
-    def __iter__(self):
+    def __await__(self):
         if not self.done():
             self._asyncio_future_blocking = True
             yield self  # This tells Task to wait for completion.
-        assert self.done(), "await wasn't used with future"
+        if not self.done():
+            raise RuntimeError("await wasn't used with future")
         return self.result()  # May raise too.
 
-    __await__ = __iter__  # make compatible with 'await' expression
+    __iter__ = __await__  # make compatible with 'yield from'.
 
 
 # Needed for testing purposes.
diff --git a/Lib/asyncio/tasks.py b/Lib/asyncio/tasks.py
index 572e707..b118088 100644
--- a/Lib/asyncio/tasks.py
+++ b/Lib/asyncio/tasks.py
@@ -37,7 +37,9 @@ def all_tasks(loop=None):
     return {t for t in _all_tasks if futures._get_loop(t) is loop}
 
 
-class Task(futures.Future):
+class Task(futures._PyFuture):  # Inherit Python Task implementation
+                                # from a Python Future implementation.
+
     """A coroutine wrapped in a Future."""
 
     # An important invariant maintained while a Task not done:
@@ -107,11 +109,17 @@ class Task(futures.Future):
             if self._source_traceback:
                 context['source_traceback'] = self._source_traceback
             self._loop.call_exception_handler(context)
-        futures.Future.__del__(self)
+        super().__del__()
 
     def _repr_info(self):
         return base_tasks._task_repr_info(self)
 
+    def set_result(self, result):
+        raise RuntimeError('Task does not support set_result operation')
+
+    def set_exception(self, exception):
+        raise RuntimeError('Task does not support set_exception operation')
+
     def get_stack(self, *, limit=None):
         """Return the list of stack frames for this task's coroutine.
 
@@ -180,7 +188,9 @@ class Task(futures.Future):
         return True
 
     def _step(self, exc=None):
-        assert not self.done(), f'_step(): already done: {self!r}, {exc!r}'
+        if self.done():
+            raise futures.InvalidStateError(
+                f'_step(): already done: {self!r}, {exc!r}')
         if self._must_cancel:
             if not isinstance(exc, futures.CancelledError):
                 exc = futures.CancelledError()
@@ -201,15 +211,15 @@ class Task(futures.Future):
             if self._must_cancel:
                 # Task is cancelled right before coro stops.
                 self._must_cancel = False
-                self.set_exception(futures.CancelledError())
+                super().set_exception(futures.CancelledError())
             else:
-                self.set_result(exc.value)
+                super().set_result(exc.value)
         except futures.CancelledError:
             super().cancel()  # I.e., Future.cancel(self).
         except Exception as exc:
-            self.set_exception(exc)
+            super().set_exception(exc)
         except BaseException as exc:
-            self.set_exception(exc)
+            super().set_exception(exc)
             raise
         else:
             blocking = getattr(result, '_asyncio_future_blocking', None)
diff --git a/Lib/test/test_asyncio/test_futures.py b/Lib/test/test_asyncio/test_futures.py
index f777a42..d339616 100644
--- a/Lib/test/test_asyncio/test_futures.py
+++ b/Lib/test/test_asyncio/test_futures.py
@@ -370,7 +370,8 @@ class BaseFutureTests:
         def test():
             arg1, arg2 = coro()
 
-        self.assertRaises(AssertionError, test)
+        with self.assertRaisesRegex(RuntimeError, "await wasn't used"):
+            test()
         fut.cancel()
 
     @mock.patch('asyncio.base_events.logger')
diff --git a/Lib/test/test_asyncio/test_tasks.py b/Lib/test/test_asyncio/test_tasks.py
index 4cf694f..0dc6c32 100644
--- a/Lib/test/test_asyncio/test_tasks.py
+++ b/Lib/test/test_asyncio/test_tasks.py
@@ -1332,17 +1332,23 @@ class BaseTaskTests:
         self.assertIsNone(task._fut_waiter)
         self.assertTrue(fut.cancelled())
 
-    def test_step_in_completed_task(self):
+    def test_task_set_methods(self):
         @asyncio.coroutine
         def notmuch():
             return 'ko'
 
         gen = notmuch()
         task = self.new_task(self.loop, gen)
-        task.set_result('ok')
 
-        self.assertRaises(AssertionError, task._step)
-        gen.close()
+        with self.assertRaisesRegex(RuntimeError, 'not support set_result'):
+            task.set_result('ok')
+
+        with self.assertRaisesRegex(RuntimeError, 'not support set_exception'):
+            task.set_exception(ValueError())
+
+        self.assertEqual(
+            self.loop.run_until_complete(task),
+            'ko')
 
     def test_step_result(self):
         @asyncio.coroutine
@@ -2231,10 +2237,59 @@ def add_subclass_tests(cls):
     return cls
 
 
+class SetMethodsTest:
+
+    def test_set_result_causes_invalid_state(self):
+        Future = type(self).Future
+        self.loop.call_exception_handler = exc_handler = mock.Mock()
+
+        async def foo():
+            await asyncio.sleep(0.1, loop=self.loop)
+            return 10
+
+        task = self.new_task(self.loop, foo())
+        Future.set_result(task, 'spam')
+
+        self.assertEqual(
+            self.loop.run_until_complete(task),
+            'spam')
+
+        exc_handler.assert_called_once()
+        exc = exc_handler.call_args[0][0]['exception']
+        with self.assertRaisesRegex(asyncio.InvalidStateError,
+                                    r'step\(\): already done'):
+            raise exc
+
+    def test_set_exception_causes_invalid_state(self):
+        class MyExc(Exception):
+            pass
+
+        Future = type(self).Future
+        self.loop.call_exception_handler = exc_handler = mock.Mock()
+
+        async def foo():
+            await asyncio.sleep(0.1, loop=self.loop)
+            return 10
+
+        task = self.new_task(self.loop, foo())
+        Future.set_exception(task, MyExc())
+
+        with self.assertRaises(MyExc):
+            self.loop.run_until_complete(task)
+
+        exc_handler.assert_called_once()
+        exc = exc_handler.call_args[0][0]['exception']
+        with self.assertRaisesRegex(asyncio.InvalidStateError,
+                                    r'step\(\): already done'):
+            raise exc
+
+
 @unittest.skipUnless(hasattr(futures, '_CFuture') and
                      hasattr(tasks, '_CTask'),
                      'requires the C _asyncio module')
-class CTask_CFuture_Tests(BaseTaskTests, test_utils.TestCase):
+class CTask_CFuture_Tests(BaseTaskTests, SetMethodsTest,
+                          test_utils.TestCase):
+
     Task = getattr(tasks, '_CTask', None)
     Future = getattr(futures, '_CFuture', None)
 
@@ -2245,11 +2300,8 @@ class CTask_CFuture_Tests(BaseTaskTests, test_utils.TestCase):
 @add_subclass_tests
 class CTask_CFuture_SubclassTests(BaseTaskTests, test_utils.TestCase):
 
-    class Task(tasks._CTask):
-        pass
-
-    class Future(futures._CFuture):
-        pass
+    Task = getattr(tasks, '_CTask', None)
+    Future = getattr(futures, '_CFuture', None)
 
 
 @unittest.skipUnless(hasattr(tasks, '_CTask'),
@@ -2257,9 +2309,7 @@ class CTask_CFuture_SubclassTests(BaseTaskTests, test_utils.TestCase):
 @add_subclass_tests
 class CTaskSubclass_PyFuture_Tests(BaseTaskTests, test_utils.TestCase):
 
-    class Task(tasks._CTask):
-        pass
-
+    Task = getattr(tasks, '_CTask', None)
     Future = futures._PyFuture
 
 
@@ -2268,15 +2318,14 @@ class CTaskSubclass_PyFuture_Tests(BaseTaskTests, test_utils.TestCase):
 @add_subclass_tests
 class PyTask_CFutureSubclass_Tests(BaseTaskTests, test_utils.TestCase):
 
-    class Future(futures._CFuture):
-        pass
-
+    Future = getattr(futures, '_CFuture', None)
     Task = tasks._PyTask
 
 
 @unittest.skipUnless(hasattr(tasks, '_CTask'),
                      'requires the C _asyncio module')
 class CTask_PyFuture_Tests(BaseTaskTests, test_utils.TestCase):
+
     Task = getattr(tasks, '_CTask', None)
     Future = futures._PyFuture
 
@@ -2284,22 +2333,22 @@ class CTask_PyFuture_Tests(BaseTaskTests, test_utils.TestCase):
 @unittest.skipUnless(hasattr(futures, '_CFuture'),
                      'requires the C _asyncio module')
 class PyTask_CFuture_Tests(BaseTaskTests, test_utils.TestCase):
+
     Task = tasks._PyTask
     Future = getattr(futures, '_CFuture', None)
 
 
-class PyTask_PyFuture_Tests(BaseTaskTests, test_utils.TestCase):
+class PyTask_PyFuture_Tests(BaseTaskTests, SetMethodsTest,
+                            test_utils.TestCase):
+
     Task = tasks._PyTask
     Future = futures._PyFuture
 
 
 @add_subclass_tests
 class PyTask_PyFuture_SubclassTests(BaseTaskTests, test_utils.TestCase):
-    class Task(tasks._PyTask):
-        pass
-
-    class Future(futures._PyFuture):
-        pass
+    Task = tasks._PyTask
+    Future = futures._PyFuture
 
 
 @unittest.skipUnless(hasattr(tasks, '_CTask'),