autoasync.py 5.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142
  1. # Copyright 2014-2015 Nathan West
  2. #
  3. # This file is part of autocommand.
  4. #
  5. # autocommand is free software: you can redistribute it and/or modify
  6. # it under the terms of the GNU Lesser General Public License as published by
  7. # the Free Software Foundation, either version 3 of the License, or
  8. # (at your option) any later version.
  9. #
  10. # autocommand is distributed in the hope that it will be useful,
  11. # but WITHOUT ANY WARRANTY; without even the implied warranty of
  12. # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
  13. # GNU Lesser General Public License for more details.
  14. #
  15. # You should have received a copy of the GNU Lesser General Public License
  16. # along with autocommand. If not, see <http://www.gnu.org/licenses/>.
  17. from asyncio import get_event_loop, iscoroutine
  18. from functools import wraps
  19. from inspect import signature
  20. async def _run_forever_coro(coro, args, kwargs, loop):
  21. '''
  22. This helper function launches an async main function that was tagged with
  23. forever=True. There are two possibilities:
  24. - The function is a normal function, which handles initializing the event
  25. loop, which is then run forever
  26. - The function is a coroutine, which needs to be scheduled in the event
  27. loop, which is then run forever
  28. - There is also the possibility that the function is a normal function
  29. wrapping a coroutine function
  30. The function is therefore called unconditionally and scheduled in the event
  31. loop if the return value is a coroutine object.
  32. The reason this is a separate function is to make absolutely sure that all
  33. the objects created are garbage collected after all is said and done; we
  34. do this to ensure that any exceptions raised in the tasks are collected
  35. ASAP.
  36. '''
  37. # Personal note: I consider this an antipattern, as it relies on the use of
  38. # unowned resources. The setup function dumps some stuff into the event
  39. # loop where it just whirls in the ether without a well defined owner or
  40. # lifetime. For this reason, there's a good chance I'll remove the
  41. # forever=True feature from autoasync at some point in the future.
  42. thing = coro(*args, **kwargs)
  43. if iscoroutine(thing):
  44. await thing
  45. def autoasync(coro=None, *, loop=None, forever=False, pass_loop=False):
  46. '''
  47. Convert an asyncio coroutine into a function which, when called, is
  48. evaluted in an event loop, and the return value returned. This is intented
  49. to make it easy to write entry points into asyncio coroutines, which
  50. otherwise need to be explictly evaluted with an event loop's
  51. run_until_complete.
  52. If `loop` is given, it is used as the event loop to run the coro in. If it
  53. is None (the default), the loop is retreived using asyncio.get_event_loop.
  54. This call is defered until the decorated function is called, so that
  55. callers can install custom event loops or event loop policies after
  56. @autoasync is applied.
  57. If `forever` is True, the loop is run forever after the decorated coroutine
  58. is finished. Use this for servers created with asyncio.start_server and the
  59. like.
  60. If `pass_loop` is True, the event loop object is passed into the coroutine
  61. as the `loop` kwarg when the wrapper function is called. In this case, the
  62. wrapper function's __signature__ is updated to remove this parameter, so
  63. that autoparse can still be used on it without generating a parameter for
  64. `loop`.
  65. This coroutine can be called with ( @autoasync(...) ) or without
  66. ( @autoasync ) arguments.
  67. Examples:
  68. @autoasync
  69. def get_file(host, port):
  70. reader, writer = yield from asyncio.open_connection(host, port)
  71. data = reader.read()
  72. sys.stdout.write(data.decode())
  73. get_file(host, port)
  74. @autoasync(forever=True, pass_loop=True)
  75. def server(host, port, loop):
  76. yield_from loop.create_server(Proto, host, port)
  77. server('localhost', 8899)
  78. '''
  79. if coro is None:
  80. return lambda c: autoasync(
  81. c, loop=loop,
  82. forever=forever,
  83. pass_loop=pass_loop)
  84. # The old and new signatures are required to correctly bind the loop
  85. # parameter in 100% of cases, even if it's a positional parameter.
  86. # NOTE: A future release will probably require the loop parameter to be
  87. # a kwonly parameter.
  88. if pass_loop:
  89. old_sig = signature(coro)
  90. new_sig = old_sig.replace(parameters=(
  91. param for name, param in old_sig.parameters.items()
  92. if name != "loop"))
  93. @wraps(coro)
  94. def autoasync_wrapper(*args, **kwargs):
  95. # Defer the call to get_event_loop so that, if a custom policy is
  96. # installed after the autoasync decorator, it is respected at call time
  97. local_loop = get_event_loop() if loop is None else loop
  98. # Inject the 'loop' argument. We have to use this signature binding to
  99. # ensure it's injected in the correct place (positional, keyword, etc)
  100. if pass_loop:
  101. bound_args = old_sig.bind_partial()
  102. bound_args.arguments.update(
  103. loop=local_loop,
  104. **new_sig.bind(*args, **kwargs).arguments)
  105. args, kwargs = bound_args.args, bound_args.kwargs
  106. if forever:
  107. local_loop.create_task(_run_forever_coro(
  108. coro, args, kwargs, local_loop
  109. ))
  110. local_loop.run_forever()
  111. else:
  112. return local_loop.run_until_complete(coro(*args, **kwargs))
  113. # Attach the updated signature. This allows 'pass_loop' to be used with
  114. # autoparse
  115. if pass_loop:
  116. autoasync_wrapper.__signature__ = new_sig
  117. return autoasync_wrapper