utils.py 4.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185
  1. import asyncio
  2. import functools
  3. import inspect
  4. import queue
  5. import random
  6. import socket
  7. import string
  8. from collections import OrderedDict
  9. from contextlib import closing
  10. from os.path import abspath, dirname
  11. import time
  12. project_dir = dirname(abspath(__file__))
  13. STATIC_PATH = '%s/html' % project_dir
  14. class ObjectDict(dict):
  15. """
  16. Object like dict, every dict[key] can visite by dict.key
  17. If dict[key] is `Get`, calculate it's value.
  18. """
  19. def __getattr__(self, name):
  20. ret = self.__getitem__(name)
  21. if hasattr(ret, '__get__'):
  22. return ret.__get__(self, ObjectDict)
  23. return ret
  24. def catch_exp_call(func, logger):
  25. """运行函数,将捕获异常记录到日志
  26. :param func: 函数
  27. :param logger: 日志
  28. :return: ``func`` 返回值
  29. """
  30. try:
  31. return func()
  32. except Exception:
  33. logger.exception("Error when invoke `%s`" % func)
  34. def iscoroutinefunction(object):
  35. while isinstance(object, functools.partial):
  36. object = object.func
  37. return asyncio.iscoroutinefunction(object)
  38. def isgeneratorfunction(object):
  39. while isinstance(object, functools.partial):
  40. object = object.func
  41. return inspect.isgeneratorfunction(object)
  42. def get_function_name(func, default=None):
  43. while isinstance(func, functools.partial):
  44. func = func.func
  45. return getattr(func, '__name__', default)
  46. class LimitedSizeQueue(queue.Queue):
  47. """
  48. 有限大小的队列
  49. `get()` 返回全部数据
  50. 队列满时,再 `put()` 会阻塞
  51. """
  52. def get(self):
  53. """获取队列全部数据"""
  54. try:
  55. return super().get(block=False)
  56. except queue.Empty:
  57. return []
  58. def wait_empty(self, timeout=None):
  59. """等待队列内的数据被取走"""
  60. with self.not_full:
  61. if self._qsize() == 0:
  62. return
  63. if timeout is None:
  64. self.not_full.wait()
  65. elif timeout < 0:
  66. raise ValueError("'timeout' must be a non-negative number")
  67. else:
  68. self.not_full.wait(timeout)
  69. def _init(self, maxsize):
  70. self.queue = []
  71. def _qsize(self):
  72. return len(self.queue)
  73. # Put a new item in the queue
  74. def _put(self, item):
  75. self.queue.append(item)
  76. # Get an item from the queue
  77. def _get(self):
  78. all_data = self.queue
  79. self.queue = []
  80. return all_data
  81. async def wait_host_port(host, port, duration=10, delay=2):
  82. """Repeatedly try if a port on a host is open until duration seconds passed
  83. from: https://gist.github.com/betrcode/0248f0fda894013382d7#gistcomment-3161499
  84. :param str host: host ip address or hostname
  85. :param int port: port number
  86. :param int/float duration: Optional. Total duration in seconds to wait, by default 10
  87. :param int/float delay: Optional. Delay in seconds between each try, by default 2
  88. :return: awaitable bool
  89. """
  90. tmax = time.time() + duration
  91. while time.time() < tmax:
  92. try:
  93. _reader, writer = await asyncio.wait_for(asyncio.open_connection(host, port), timeout=5)
  94. writer.close()
  95. await writer.wait_closed()
  96. return True
  97. except Exception:
  98. if delay:
  99. await asyncio.sleep(delay)
  100. return False
  101. def get_free_port():
  102. """
  103. pick a free port number
  104. :return int: port number
  105. """
  106. with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as s:
  107. s.bind(('', 0))
  108. s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
  109. return s.getsockname()[1]
  110. def random_str(len=16):
  111. """生成小写字母和数组组成的随机字符串
  112. :param int len: 字符串长度
  113. """
  114. return ''.join(random.SystemRandom().choice(string.ascii_lowercase + string.digits) for _ in range(len))
  115. def run_as_function(gen):
  116. res = None
  117. while 1:
  118. try:
  119. res = gen.send(res)
  120. except StopIteration as e:
  121. if len(e.args) == 1:
  122. return e.args[0]
  123. return
  124. async def to_coroutine(gen):
  125. res = None
  126. while 1:
  127. try:
  128. c = gen.send(res)
  129. res = await c
  130. except StopIteration as e:
  131. if len(e.args) == 1:
  132. return e.args[0]
  133. return
  134. class LRUDict(OrderedDict):
  135. """
  136. Store items in the order the keys were last recent updated.
  137. The last recent updated item was in end.
  138. The last furthest updated item was in front.
  139. """
  140. def __setitem__(self, key, value):
  141. OrderedDict.__setitem__(self, key, value)
  142. self.move_to_end(key)