1__all__ = ['coroutine',
2           'iscoroutinefunction', 'iscoroutine',
3           'From', 'Return']
4
5import functools
6import inspect
7import opcode
8import os
9import sys
10import textwrap
11import traceback
12import types
13
14from . import compat
15from . import events
16from . import futures
17from .log import logger
18
19
20# Opcode of "yield from" instruction
21_YIELD_FROM = opcode.opmap.get('YIELD_FROM', None)
22
23# If you set _DEBUG to true, @coroutine will wrap the resulting
24# generator objects in a CoroWrapper instance (defined below).  That
25# instance will log a message when the generator is never iterated
26# over, which may happen when you forget to use "yield from" with a
27# coroutine call.  Note that the value of the _DEBUG flag is taken
28# when the decorator is used, so to be of any use it must be set
29# before you define your coroutines.  A downside of using this feature
30# is that tracebacks show entries for the CoroWrapper.__next__ method
31# when _DEBUG is true.
32_DEBUG = bool(os.environ.get('TROLLIUSDEBUG'))
33
34
35try:
36    _types_coroutine = types.coroutine
37except AttributeError:
38    _types_coroutine = None
39
40try:
41    _inspect_iscoroutinefunction = inspect.iscoroutinefunction
42except AttributeError:
43    _inspect_iscoroutinefunction = lambda func: False
44
45try:
46    from collections.abc import Coroutine as _CoroutineABC, \
47                                Awaitable as _AwaitableABC
48except ImportError:
49    _CoroutineABC = _AwaitableABC = None
50
51
52if _YIELD_FROM is not None:
53    # Check for CPython issue #21209
54    exec('''if 1:
55        def has_yield_from_bug():
56            class MyGen:
57                def __init__(self):
58                    self.send_args = None
59                def __iter__(self):
60                    return self
61                def __next__(self):
62                    return 42
63                def send(self, *what):
64                    self.send_args = what
65                    return None
66            def yield_from_gen(gen):
67                yield from gen
68            value = (1, 2, 3)
69            gen = MyGen()
70            coro = yield_from_gen(gen)
71            next(coro)
72            coro.send(value)
73            return gen.send_args != (value,)
74''')
75    _YIELD_FROM_BUG = has_yield_from_bug()
76    del has_yield_from_bug
77else:
78    _YIELD_FROM_BUG = False
79
80
81if compat.PY35:
82    return_base_class = Exception
83else:
84    return_base_class = StopIteration
85
86class ReturnException(return_base_class):
87    def __init__(self, *args):
88        return_base_class.__init__(self)
89        if not args:
90            self.value = None
91        elif len(args) == 1:
92            self.value = args[0]
93        else:
94            self.value = args
95        self.raised = False
96        if _DEBUG:
97            frame = sys._getframe(1)
98            self._source_traceback = traceback.extract_stack(frame)
99            # explicitly clear the reference to avoid reference cycles
100            frame = None
101        else:
102            self._source_traceback = None
103
104    def __del__(self):
105        if self.raised:
106            return
107
108        fmt = 'Return(%r) used without raise'
109        if self._source_traceback:
110            fmt += '\nReturn created at (most recent call last):\n'
111            tb = ''.join(traceback.format_list(self._source_traceback))
112            fmt += tb.rstrip()
113        logger.error(fmt, self.value)
114
115
116if compat.PY33 and not compat.PY35:
117    # Don't use the Return class on Python 3.3 and 3.4 to support asyncio
118    # coroutines (to avoid the warning emited in Return destructor).
119    #
120    # The problem is that ReturnException inherits from StopIteration.
121    # "yield from trollius_coroutine". Task._step() does not receive the Return
122    # exception, because "yield from" handles it internally. So it's not
123    # possible to set the raised attribute to True to avoid the warning in
124    # Return destructor.
125    def Return(*args):
126        if not args:
127            value = None
128        elif len(args) == 1:
129            value = args[0]
130        else:
131            value = args
132        return StopIteration(value)
133else:
134    Return = ReturnException
135
136
137def debug_wrapper(gen):
138    # This function is called from 'sys.set_coroutine_wrapper'.
139    # We only wrap here coroutines defined via 'async def' syntax.
140    # Generator-based coroutines are wrapped in @coroutine
141    # decorator.
142    return CoroWrapper(gen, None)
143
144
145def _coroutine_at_yield_from(coro):
146    """Test if the last instruction of a coroutine is "yield from".
147
148    Return False if the coroutine completed.
149    """
150    frame = coro.gi_frame
151    if frame is None:
152        return False
153    code = coro.gi_code
154    assert frame.f_lasti >= 0
155    offset = frame.f_lasti + 1
156    instr = code.co_code[offset]
157    return (instr == _YIELD_FROM)
158
159
160class CoroWrapper:
161    # Wrapper for coroutine object in _DEBUG mode.
162
163    def __init__(self, gen, func=None):
164        assert inspect.isgenerator(gen) or inspect.iscoroutine(gen), gen
165        self.gen = gen
166        self.func = func # Used to unwrap @coroutine decorator
167        self._source_traceback = traceback.extract_stack(sys._getframe(1))
168        self.__name__ = getattr(gen, '__name__', None)
169        self.__qualname__ = getattr(gen, '__qualname__', None)
170
171    def __repr__(self):
172        coro_repr = _format_coroutine(self)
173        if self._source_traceback:
174            frame = self._source_traceback[-1]
175            coro_repr += ', created at %s:%s' % (frame[0], frame[1])
176        return '<%s %s>' % (self.__class__.__name__, coro_repr)
177
178    def __iter__(self):
179        return self
180
181    def __next__(self):
182        return next(self.gen)
183    next = __next__
184
185    if _YIELD_FROM_BUG:
186        # For for CPython issue #21209: using "yield from" and a custom
187        # generator, generator.send(tuple) unpacks the tuple instead of passing
188        # the tuple unchanged. Check if the caller is a generator using "yield
189        # from" to decide if the parameter should be unpacked or not.
190        def send(self, *value):
191            frame = sys._getframe()
192            caller = frame.f_back
193            assert caller.f_lasti >= 0
194            if caller.f_code.co_code[caller.f_lasti] != _YIELD_FROM:
195                value = value[0]
196            return self.gen.send(value)
197    else:
198        def send(self, value):
199            return self.gen.send(value)
200
201    def throw(self, exc_type, exc_value=None, exc_tb=None):
202        return self.gen.throw(exc_type, exc_value, exc_tb)
203
204    def close(self):
205        return self.gen.close()
206
207    @property
208    def gi_frame(self):
209        return self.gen.gi_frame
210
211    @property
212    def gi_running(self):
213        return self.gen.gi_running
214
215    @property
216    def gi_code(self):
217        return self.gen.gi_code
218
219    if compat.PY35:
220
221        __await__ = __iter__ # make compatible with 'await' expression
222
223        @property
224        def gi_yieldfrom(self):
225            return self.gen.gi_yieldfrom
226
227        @property
228        def cr_await(self):
229            return self.gen.cr_await
230
231        @property
232        def cr_running(self):
233            return self.gen.cr_running
234
235        @property
236        def cr_code(self):
237            return self.gen.cr_code
238
239        @property
240        def cr_frame(self):
241            return self.gen.cr_frame
242
243    def __del__(self):
244        # Be careful accessing self.gen.frame -- self.gen might not exist.
245        gen = getattr(self, 'gen', None)
246        frame = getattr(gen, 'gi_frame', None)
247        if frame is None:
248            frame = getattr(gen, 'cr_frame', None)
249        if frame is not None and frame.f_lasti == -1:
250            msg = '%r was never yielded from' % self
251            tb = getattr(self, '_source_traceback', ())
252            if tb:
253                tb = ''.join(traceback.format_list(tb))
254                msg += ('\nCoroutine object created at '
255                        '(most recent call last):\n')
256                msg += tb.rstrip()
257            logger.error(msg)
258
259if not compat.PY34:
260    # Backport functools.update_wrapper() from Python 3.4:
261    # - Python 2.7 fails if assigned attributes don't exist
262    # - Python 2.7 and 3.1 don't set the __wrapped__ attribute
263    # - Python 3.2 and 3.3 set __wrapped__ before updating __dict__
264    def _update_wrapper(wrapper,
265                       wrapped,
266                       assigned = functools.WRAPPER_ASSIGNMENTS,
267                       updated = functools.WRAPPER_UPDATES):
268        """Update a wrapper function to look like the wrapped function
269
270           wrapper is the function to be updated
271           wrapped is the original function
272           assigned is a tuple naming the attributes assigned directly
273           from the wrapped function to the wrapper function (defaults to
274           functools.WRAPPER_ASSIGNMENTS)
275           updated is a tuple naming the attributes of the wrapper that
276           are updated with the corresponding attribute from the wrapped
277           function (defaults to functools.WRAPPER_UPDATES)
278        """
279        for attr in assigned:
280            try:
281                value = getattr(wrapped, attr)
282            except AttributeError:
283                pass
284            else:
285                setattr(wrapper, attr, value)
286        for attr in updated:
287            getattr(wrapper, attr).update(getattr(wrapped, attr, {}))
288        # Issue #17482: set __wrapped__ last so we don't inadvertently copy it
289        # from the wrapped function when updating __dict__
290        wrapper.__wrapped__ = wrapped
291        # Return the wrapper so this can be used as a decorator via partial()
292        return wrapper
293
294    def _wraps(wrapped,
295              assigned = functools.WRAPPER_ASSIGNMENTS,
296              updated = functools.WRAPPER_UPDATES):
297        """Decorator factory to apply update_wrapper() to a wrapper function
298
299           Returns a decorator that invokes update_wrapper() with the decorated
300           function as the wrapper argument and the arguments to wraps() as the
301           remaining arguments. Default arguments are as for update_wrapper().
302           This is a convenience function to simplify applying partial() to
303           update_wrapper().
304        """
305        return functools.partial(_update_wrapper, wrapped=wrapped,
306                                 assigned=assigned, updated=updated)
307else:
308    _wraps = functools.wraps
309
310_PEP479 = (sys.version_info >= (3, 5))
311if _PEP479:
312    # Need exec() because yield+return raises a SyntaxError on Python 2
313    exec(textwrap.dedent('''
314        def pep479_wrapper(func, coro_func):
315            @_wraps(func)
316            def pep479_wrapped(*args, **kw):
317                coro = coro_func(*args, **kw)
318                value = None
319                error = None
320                while True:
321                    try:
322                        if error is not None:
323                            value = coro.throw(error)
324                        elif value is not None:
325                            value = coro.send(value)
326                        else:
327                            value = next(coro)
328                    except RuntimeError:
329                        # FIXME: special case for
330                        # FIXME: "isinstance(exc.__context__, StopIteration)"?
331                        raise
332                    except StopIteration as exc:
333                        return exc.value
334                    except Return as exc:
335                        exc.raised = True
336                        return exc.value
337                    except BaseException as exc:
338                        raise
339
340                    try:
341                        value = yield value
342                        error = None
343                    except BaseException as exc:
344                        value = None
345                        error = exc
346
347            return pep479_wrapped
348        '''))
349
350
351def coroutine(func):
352    """Decorator to mark coroutines.
353
354    If the coroutine is not yielded from before it is destroyed,
355    an error message is logged.
356    """
357    if _inspect_iscoroutinefunction(func):
358        # In Python 3.5 that's all we need to do for coroutines
359        # defiend with "async def".
360        # Wrapping in CoroWrapper will happen via
361        # 'sys.set_coroutine_wrapper' function.
362        return func
363
364    if inspect.isgeneratorfunction(func):
365        coro = func
366    else:
367        @_wraps(func)
368        def coro(*args, **kw):
369            res = func(*args, **kw)
370            if (isinstance(res, futures._FUTURE_CLASSES)
371                or inspect.isgenerator(res)):
372                res = yield From(res)
373            elif _AwaitableABC is not None:
374                # If 'func' returns an Awaitable (new in 3.5) we
375                # want to run it.
376                try:
377                    await_meth = res.__await__
378                except AttributeError:
379                    pass
380                else:
381                    if isinstance(res, _AwaitableABC):
382                        res = yield From(await_meth())
383            raise Return(res)
384
385    if _PEP479:
386        # FIXME: use @_wraps
387        coro = pep479_wrapper(func, coro)
388        coro = _wraps(func)(coro)
389
390    if not _DEBUG:
391        if _types_coroutine is None:
392            wrapper = coro
393        else:
394            wrapper = _types_coroutine(coro)
395    else:
396        @_wraps(func)
397        def wrapper(*args, **kwds):
398            w = CoroWrapper(coro(*args, **kwds), func=func)
399            if w._source_traceback:
400                del w._source_traceback[-1]
401            # Python < 3.5 does not implement __qualname__
402            # on generator objects, so we set it manually.
403            # We use getattr as some callables (such as
404            # functools.partial may lack __qualname__).
405            w.__name__ = getattr(func, '__name__', None)
406            w.__qualname__ = getattr(func, '__qualname__', None)
407            return w
408
409    wrapper._is_coroutine = True  # For iscoroutinefunction().
410    return wrapper
411
412
413def iscoroutinefunction(func):
414    """Return True if func is a decorated coroutine function."""
415    return (getattr(func, '_is_coroutine', False) or
416            _inspect_iscoroutinefunction(func))
417
418
419_COROUTINE_TYPES = (types.GeneratorType, CoroWrapper)
420if _CoroutineABC is not None:
421    _COROUTINE_TYPES += (_CoroutineABC,)
422if events.asyncio is not None:
423    # Accept also asyncio CoroWrapper for interoperability
424    if hasattr(events.asyncio, 'coroutines'):
425        _COROUTINE_TYPES += (events.asyncio.coroutines.CoroWrapper,)
426    else:
427        # old asyncio/Python versions
428        _COROUTINE_TYPES += (events.asyncio.tasks.CoroWrapper,)
429
430def iscoroutine(obj):
431    """Return True if obj is a coroutine object."""
432    return isinstance(obj, _COROUTINE_TYPES)
433
434
435def _format_coroutine(coro):
436    assert iscoroutine(coro)
437
438    coro_name = None
439    if isinstance(coro, CoroWrapper):
440        func = coro.func
441        coro_name = coro.__qualname__
442        if coro_name is not None:
443            coro_name = '{0}()'.format(coro_name)
444    else:
445        func = coro
446
447    if coro_name is None:
448        coro_name = events._format_callback(func, ())
449
450    try:
451        coro_code = coro.gi_code
452    except AttributeError:
453        coro_code = coro.cr_code
454
455    try:
456        coro_frame = coro.gi_frame
457    except AttributeError:
458        coro_frame = coro.cr_frame
459
460    filename = coro_code.co_filename
461    if (isinstance(coro, CoroWrapper)
462    and not inspect.isgeneratorfunction(coro.func)
463    and coro.func is not None):
464        filename, lineno = events._get_function_source(coro.func)
465        if coro_frame is None:
466            coro_repr = ('%s done, defined at %s:%s'
467                         % (coro_name, filename, lineno))
468        else:
469            coro_repr = ('%s running, defined at %s:%s'
470                         % (coro_name, filename, lineno))
471    elif coro_frame is not None:
472        lineno = coro_frame.f_lineno
473        coro_repr = ('%s running at %s:%s'
474                     % (coro_name, filename, lineno))
475    else:
476        lineno = coro_code.co_firstlineno
477        coro_repr = ('%s done, defined at %s:%s'
478                     % (coro_name, filename, lineno))
479
480    return coro_repr
481
482
483class FromWrapper(object):
484    __slots__ = ('obj',)
485
486    def __init__(self, obj):
487        if isinstance(obj, FromWrapper):
488            obj = obj.obj
489            assert not isinstance(obj, FromWrapper)
490        self.obj = obj
491
492def From(obj):
493    if not _DEBUG:
494        return obj
495    else:
496        return FromWrapper(obj)
497