Support top level await (#158)

* Support top level await

* Fix CI installing wrong wheel

* Revert "Fix CI installing wrong wheel"

This reverts commit bbf6952.

* Add documentation for top level awaits

* Update changelog

* Minor linting

* Break out error handling

* Rename exception class to ExistingEventLoopError

---------

Co-authored-by: joncrall <erotemic@gmail.com>

CHANGELOG.md

@@ -7,6 +7,8 @@ This project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.htm
 
 ## Version 1.2.0 - Unreleased
 
+### Added
+* Support for top level awaits in async code examples.
 
 ### Removed
 * Dropped 3.6 and 3.7 support. Now supporting 3.6+ Use xdoctest<=1.1.6 for 3.6 or 3.7 support.

README.rst

@@ -328,6 +328,7 @@ The main enhancements ``xdoctest`` offers over ``doctest`` are:
    stdout, both are checked, and the test will pass if either matches.
 6. Ouptut from multiple sequential print statements can now be checked by
    a single "got" statement. (new in 0.4.0).
+7. Examples can include `async code at the top level <https://xdoctest.readthedocs.io/en/latest/manual/async_doctest.html>`__.
 
 See code in ``dev/_compare/demo_enhancements.py`` for a demo that illustrates
 several of these enhancements. This demo shows cases where ``xdoctest`` works

docs/source/index.rst

@@ -6,9 +6,10 @@
 
 .. toctree::
    :maxdepth: 8
-   :caption: Package Layout
 
-   auto/xdoctest
+   Package layout <auto/xdoctest>
+   manual/xdoc_with_jupyter
+   manual/async_doctest
 
 
 Indices and tables

docs/source/manual/async_doctest.rst

@@ -0,0 +1,59 @@
+Doctests with async code
+------------------------
+
+Python 3.5 introduced `async` functions.  These are functions that run within an "event loop" that is not blocked if those functions perform blocking IO. It's similar to writing multi-threaded code but with the advantage of not having to worry about thread safety.  For more information see `the python docs <https://peps.python.org/pep-0492/>`__.
+
+Asynchronous python code examples using `asyncio <https://docs.python.org/3/library/asyncio.html>`__ are supported at the top level by xdoctest.
+This means that your code examples do not have to wrap every small snippet in a function and call `asyncio.run`.
+xdoctest handles that for you keeping the examples simple and easy to follow.
+
+For example **without xdoctest** your code example would have to be written like this:
+
+.. code:: python
+
+    >>> import yourlibrary
+    >>> import asyncio
+    >>> async def connect_and_get_running_info_wrapper():
+    ...     server = await yourlibrary.connect_to_server("example_server_url")
+    ...     running_info = await server.get_running_info()
+    ...     return server, running_info
+    ...
+    >>> server, running_info = asyncio.run(connect_and_get_running_info_wrapper())
+    >>> running_info.restarted_at
+    01:00
+
+    >>> async def restart_and_get_running_info_wrapper(server):
+    ...     await server.restart()
+    ...     return await server.get_running_info()
+    ...
+    >>> running_info = asyncio.run(restart_and_get_running_info_wrapper(server))
+    >>> running_info.restarted_at
+    13:15
+
+Now **with xdoctest** this can now be written like this:
+
+.. code:: python
+
+    >>> import yourlibrary
+    >>> server = await yourlibrary.connect_to_server("example_server_url")
+    >>> running_info = await server.get_running_info()
+    >>> running_info.restarted_at
+    01:00
+
+    >>> await server.restart()
+    >>> running_info = await server.get_running_info()
+    >>> running_info.restarted_at
+    13:15
+
+The improvement in brevity is obvious but even more so if you are writing longer examples where you want to maintain and reuse variables between each test output step.
+
+.. note::
+
+    If you don't want to utilise this feature for your async examples you don't have to.  Just don't write code examples with top level awaits.
+
+Caveats
+=======
+
+* Consumers reading your documentation may not be familiar with async concepts. It could be helpful to mention in your docs that the code examples should be run in an event loop or in a REPL that supports top-level ``await``. (IPython supports this by default. For the standard Python REPL, use ``python -m asyncio``.)
+* Using top level awaits in tests that are already running in an event loop is not supported.
+* Only python's native asyncio library is supported for top level awaits.

src/xdoctest/__init__.py

@@ -326,8 +326,9 @@ def fib(n):
 from xdoctest import docstr
 from xdoctest.runner import (doctest_module, doctest_callable,)
 from xdoctest.exceptions import (DoctestParseError, ExitTestException,
-                                 MalformedDocstr,)
+                                 MalformedDocstr, ExistingEventLoopError)
 
 __all__ = ['DoctestParseError', 'ExitTestException', 'MalformedDocstr',
+           'ExistingEventLoopError',
            'doctest_module', 'doctest_callable', 'utils', 'docstr',
            '__version__']

src/xdoctest/doctest_example.py

@@ -2,12 +2,16 @@
 This module defines the main class that holds a DocTest example
 """
 import __future__
+import asyncio
+import ast
 from collections import OrderedDict
 import traceback
 import warnings
 import math
 import sys
 import re
+import types
+from inspect import CO_COROUTINE
 from xdoctest import utils
 from xdoctest import directive
 from xdoctest import constants
@@ -272,7 +276,6 @@ def __init__(self, docsrc, modpath=None, callname=None, num=0,
             block_type (str | None):
             mode (str):
         """
-        import types
         # if we know the google block type it is recorded
         self.block_type = block_type
 
@@ -668,6 +671,7 @@ def _test_globals(self):
         # force print function and division futures
         compileflags |= __future__.print_function.compiler_flag
         compileflags |= __future__.division.compiler_flag
+        compileflags |= ast.PyCF_ALLOW_TOP_LEVEL_AWAIT
         return test_globals, compileflags
 
     def anything_ran(self):
@@ -721,6 +725,11 @@ def run(self, verbose=None, on_error=None):
         #     runstate['SKIP'] = True
 
         needs_capture = True
+        try:
+            asyncio.get_running_loop()
+            is_running_in_loop = True
+        except RuntimeError:
+            is_running_in_loop = False
 
         DEBUG = global_state.DEBUG_DOCTEST
 
@@ -843,7 +852,17 @@ def run(self, verbose=None, on_error=None):
                             # expect it to return an object with a repr that
                             # can compared to a "want" statement.
                             # print('part.compile_mode = {!r}'.format(part.compile_mode))
-                            if part.compile_mode == 'eval':
+                            if code.co_flags & CO_COROUTINE == CO_COROUTINE:
+                                if is_running_in_loop:
+                                    raise exceptions.ExistingEventLoopError(
+                                        "Cannot run top-level await doctests from within a running event loop: %s",
+                                        part.orig_lines
+                                        )
+                                if part.compile_mode == 'eval':
+                                    got_eval = asyncio.run(eval(code, test_globals))
+                                else:
+                                    asyncio.run(eval(code, test_globals))
+                            elif part.compile_mode == 'eval':
                                 got_eval = eval(code, test_globals)
                             else:
                                 exec(code, test_globals)
@@ -890,6 +909,13 @@ def run(self, verbose=None, on_error=None):
                     if verbose > 0:
                         print('Test gracefully exists on: ex={}'.format(ex))
                     break
+                except exceptions.ExistingEventLoopError:
+                    # When we try to run a doctest with await, but there is
+                    # already a running event loop.
+                    self.exc_info = sys.exc_info()
+                    if on_error == 'raise':
+                        raise
+                    break
                 except checker.GotWantException:
                     # When the "got", doesn't match the "want"
                     self.exc_info = sys.exc_info()
@@ -1038,7 +1064,7 @@ def failed_line_offset(self):
                 return 0
             ex_type, ex_value, tb = self.exc_info
             offset = self.failed_part.line_offset
-            if isinstance(ex_value, checker.ExtractGotReprException):
+            if isinstance(ex_value, (checker.ExtractGotReprException, exceptions.ExistingEventLoopError)):
                 # Return the line of the "got" expression
                 offset += self.failed_part.n_exec_lines
             elif isinstance(ex_value, checker.GotWantException):

src/xdoctest/exceptions.py

@@ -8,7 +8,13 @@ class MalformedDocstr(Exception):
     Exception raised when the docstring itself does not conform to the expected
     style (e.g. google / numpy).
     """
-    pass
+
+
+class ExistingEventLoopError(Exception):
+    """
+    Exception raised when the docstring uses a top level await and the test is
+    already running in an event loop.
+    """
 
 
 class DoctestParseError(Exception):
@@ -38,7 +44,6 @@ class IncompleteParseError(SyntaxError):
     """
     Used when something goes wrong in the xdoctest parser
     """
-    pass
 
 try:
     import _pytest

src/xdoctest/exceptions.pyi

@@ -4,6 +4,8 @@ from typing import Any
 class MalformedDocstr(Exception):
     ...
 
+class ExistingEventLoopError(Exception):
+    ...
 
 class DoctestParseError(Exception):
     msg: str

tests/test_doctest_example.py

@@ -1,4 +1,5 @@
 from xdoctest import doctest_example
+from xdoctest import exceptions
 from xdoctest import utils
 from xdoctest import constants
 from xdoctest import checker
@@ -262,6 +263,113 @@ def test_want_error_msg_failure():
         self.run(on_error='raise')
 
 
+def test_await():
+    """
+    python tests/test_doctest_example.py test_await
+    pytest tests/test_doctest_example.py::test_await
+    """
+    string = utils.codeblock(
+        '''
+        >>> import asyncio
+        >>> res = await asyncio.sleep(0, result="slept")
+        >>> print(res)
+        slept
+        ''')
+    self = doctest_example.DocTest(docsrc=string)
+    result = self.run(on_error='raise')
+    assert result['passed']
+
+
+def test_async_for():
+    """
+    python tests/test_doctest_example.py test_await
+    pytest tests/test_doctest_example.py::test_await
+    """
+    string = utils.codeblock(
+        '''
+        >>> async def test_gen():
+        >>>     yield 5
+        >>>     yield 6
+        >>> async for i in test_gen():
+        >>>     print(i)
+        5
+        6
+        ''')
+    self = doctest_example.DocTest(docsrc=string)
+    result = self.run(on_error='raise')
+    assert result['passed']
+
+
+def test_async_with():
+    """
+    python tests/test_doctest_example.py test_await
+    pytest tests/test_doctest_example.py::test_await
+    """
+    string = utils.codeblock(
+        '''
+        >>> from contextlib import asynccontextmanager
+        >>> import asyncio
+        >>> @asynccontextmanager
+        >>> async def gen():
+        >>>     try:
+        >>>         yield 1
+        >>>     finally:
+        >>>         await asyncio.sleep(0)
+        >>> async with gen() as res:
+        >>>     print(res)
+        1
+        ''')
+    self = doctest_example.DocTest(docsrc=string)
+    result = self.run(on_error='raise')
+    assert result['passed']
+
+
+def test_await_in_running_loop():
+    """
+    python tests/test_doctest_example.py test_await
+    pytest tests/test_doctest_example.py::test_await
+    """
+    string = utils.codeblock(
+        '''
+        >>> import asyncio
+        >>> res = await asyncio.sleep(0, result="slept")
+        >>> print(res)
+        slept
+        ''')
+    import asyncio
+    import pytest
+    self = doctest_example.DocTest(docsrc=string)
+    async def run_in_loop(doctest, on_error, verbose=None):
+        return doctest.run(on_error=on_error, verbose=verbose)
+
+    with pytest.raises(exceptions.ExistingEventLoopError):
+        asyncio.run(run_in_loop(self, 'raise'))
+
+    self = doctest_example.DocTest(docsrc=string)
+    res = asyncio.run(run_in_loop(self, 'return', verbose=3))
+
+    assert res['failed']
+    assert self.repr_failure()
+
+
+def test_async_def():
+    """
+    python tests/test_doctest_example.py test_await
+    pytest tests/test_doctest_example.py::test_await
+    """
+    string = utils.codeblock(
+        '''
+        >>> import asyncio
+        >>> async def run_async():
+        >>>     return await asyncio.sleep(0, result="slept")
+        >>> res = asyncio.run(run_async())
+        >>> print(res)
+        slept
+        ''')
+    self = doctest_example.DocTest(docsrc=string)
+    result = self.run(on_error='raise')
+    assert result['passed']
+
 if __name__ == '__main__':
     """
     CommandLine: