/asyncpg/compat.py

https://github.com/MagicStack/asyncpg · Python · 108 lines · 75 code · 21 blank · 12 comment · 23 complexity · b561ab5c50ddde3fcb49d72c0efab728 MD5 · raw file

  1. # Copyright (C) 2016-present the asyncpg authors and contributors
  2. # <see AUTHORS file>
  3. #
  4. # This module is part of asyncpg and is released under
  5. # the Apache 2.0 License: http://www.apache.org/licenses/LICENSE-2.0
  6. import asyncio
  7. import functools
  8. import os
  9. import pathlib
  10. import platform
  11. import sys
  12. PY_36 = sys.version_info >= (3, 6)
  13. PY_37 = sys.version_info >= (3, 7)
  14. SYSTEM = platform.uname().system
  15. if sys.version_info < (3, 5, 2):
  16. def aiter_compat(func):
  17. @functools.wraps(func)
  18. async def wrapper(self):
  19. return func(self)
  20. return wrapper
  21. else:
  22. def aiter_compat(func):
  23. return func
  24. if PY_36:
  25. fspath = os.fspath
  26. else:
  27. def fspath(path):
  28. fsp = getattr(path, '__fspath__', None)
  29. if fsp is not None and callable(fsp):
  30. path = fsp()
  31. if not isinstance(path, (str, bytes)):
  32. raise TypeError(
  33. 'expected {}() to return str or bytes, not {}'.format(
  34. fsp.__qualname__, type(path).__name__
  35. ))
  36. return path
  37. elif isinstance(path, (str, bytes)):
  38. return path
  39. else:
  40. raise TypeError(
  41. 'expected str, bytes or path-like object, not {}'.format(
  42. type(path).__name__
  43. )
  44. )
  45. if SYSTEM == 'Windows':
  46. import ctypes.wintypes
  47. CSIDL_APPDATA = 0x001a
  48. def get_pg_home_directory() -> pathlib.Path:
  49. # We cannot simply use expanduser() as that returns the user's
  50. # home directory, whereas Postgres stores its config in
  51. # %AppData% on Windows.
  52. buf = ctypes.create_unicode_buffer(ctypes.wintypes.MAX_PATH)
  53. r = ctypes.windll.shell32.SHGetFolderPathW(0, CSIDL_APPDATA, 0, 0, buf)
  54. if r:
  55. return None
  56. else:
  57. return pathlib.Path(buf.value) / 'postgresql'
  58. else:
  59. def get_pg_home_directory() -> pathlib.Path:
  60. return pathlib.Path.home()
  61. if PY_37:
  62. def current_asyncio_task(loop):
  63. return asyncio.current_task(loop)
  64. else:
  65. def current_asyncio_task(loop):
  66. return asyncio.Task.current_task(loop)
  67. async def wait_closed(stream):
  68. # Not all asyncio versions have StreamWriter.wait_closed().
  69. if hasattr(stream, 'wait_closed'):
  70. try:
  71. await stream.wait_closed()
  72. except ConnectionResetError:
  73. # On Windows wait_closed() sometimes propagates
  74. # ConnectionResetError which is totally unnecessary.
  75. pass
  76. # Workaround for https://bugs.python.org/issue37658
  77. async def wait_for(fut, timeout):
  78. if timeout is None:
  79. return await fut
  80. fut = asyncio.ensure_future(fut)
  81. try:
  82. return await asyncio.wait_for(fut, timeout)
  83. except asyncio.CancelledError:
  84. if fut.done():
  85. return fut.result()
  86. else:
  87. raise