__init__.py 3.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139
  1. """Handle files using a thread pool executor."""
  2. import asyncio
  3. import sys
  4. from functools import partial, singledispatch
  5. from io import (
  6. BufferedIOBase,
  7. BufferedRandom,
  8. BufferedReader,
  9. BufferedWriter,
  10. FileIO,
  11. TextIOBase,
  12. )
  13. from ..base import AiofilesContextManager
  14. from .binary import (
  15. AsyncBufferedIOBase,
  16. AsyncBufferedReader,
  17. AsyncFileIO,
  18. AsyncIndirectBufferedIOBase,
  19. )
  20. from .text import AsyncTextIndirectIOWrapper, AsyncTextIOWrapper
  21. sync_open = open
  22. __all__ = (
  23. "open",
  24. "stdin",
  25. "stdout",
  26. "stderr",
  27. "stdin_bytes",
  28. "stdout_bytes",
  29. "stderr_bytes",
  30. )
  31. def open(
  32. file,
  33. mode="r",
  34. buffering=-1,
  35. encoding=None,
  36. errors=None,
  37. newline=None,
  38. closefd=True,
  39. opener=None,
  40. *,
  41. loop=None,
  42. executor=None,
  43. ):
  44. return AiofilesContextManager(
  45. _open(
  46. file,
  47. mode=mode,
  48. buffering=buffering,
  49. encoding=encoding,
  50. errors=errors,
  51. newline=newline,
  52. closefd=closefd,
  53. opener=opener,
  54. loop=loop,
  55. executor=executor,
  56. )
  57. )
  58. async def _open(
  59. file,
  60. mode="r",
  61. buffering=-1,
  62. encoding=None,
  63. errors=None,
  64. newline=None,
  65. closefd=True,
  66. opener=None,
  67. *,
  68. loop=None,
  69. executor=None,
  70. ):
  71. """Open an asyncio file."""
  72. if loop is None:
  73. loop = asyncio.get_running_loop()
  74. cb = partial(
  75. sync_open,
  76. file,
  77. mode=mode,
  78. buffering=buffering,
  79. encoding=encoding,
  80. errors=errors,
  81. newline=newline,
  82. closefd=closefd,
  83. opener=opener,
  84. )
  85. f = await loop.run_in_executor(executor, cb)
  86. return wrap(f, loop=loop, executor=executor)
  87. @singledispatch
  88. def wrap(file, *, loop=None, executor=None):
  89. raise TypeError("Unsupported io type: {}.".format(file))
  90. @wrap.register(TextIOBase)
  91. def _(file, *, loop=None, executor=None):
  92. return AsyncTextIOWrapper(file, loop=loop, executor=executor)
  93. @wrap.register(BufferedWriter)
  94. @wrap.register(BufferedIOBase)
  95. def _(file, *, loop=None, executor=None):
  96. return AsyncBufferedIOBase(file, loop=loop, executor=executor)
  97. @wrap.register(BufferedReader)
  98. @wrap.register(BufferedRandom)
  99. def _(file, *, loop=None, executor=None):
  100. return AsyncBufferedReader(file, loop=loop, executor=executor)
  101. @wrap.register(FileIO)
  102. def _(file, *, loop=None, executor=None):
  103. return AsyncFileIO(file, loop=loop, executor=executor)
  104. stdin = AsyncTextIndirectIOWrapper("sys.stdin", None, None, indirect=lambda: sys.stdin)
  105. stdout = AsyncTextIndirectIOWrapper(
  106. "sys.stdout", None, None, indirect=lambda: sys.stdout
  107. )
  108. stderr = AsyncTextIndirectIOWrapper(
  109. "sys.stderr", None, None, indirect=lambda: sys.stderr
  110. )
  111. stdin_bytes = AsyncIndirectBufferedIOBase(
  112. "sys.stdin.buffer", None, None, indirect=lambda: sys.stdin.buffer
  113. )
  114. stdout_bytes = AsyncIndirectBufferedIOBase(
  115. "sys.stdout.buffer", None, None, indirect=lambda: sys.stdout.buffer
  116. )
  117. stderr_bytes = AsyncIndirectBufferedIOBase(
  118. "sys.stderr.buffer", None, None, indirect=lambda: sys.stderr.buffer
  119. )