PageRenderTime 45ms CodeModel.GetById 11ms RepoModel.GetById 0ms app.codeStats 0ms

/homeassistant/util/__init__.py

https://gitlab.com/mkerfoot/home-assistant
Python | 415 lines | 404 code | 4 blank | 7 comment | 0 complexity | 298739285665c6f3a4cb7cc9c4c9fb48 MD5 | raw file
  1. """
  2. homeassistant.util
  3. ~~~~~~~~~~~~~~~~~~
  4. Helper methods for various modules.
  5. """
  6. import collections
  7. from itertools import chain
  8. import threading
  9. import queue
  10. from datetime import datetime
  11. import re
  12. import enum
  13. import socket
  14. import random
  15. import string
  16. from functools import wraps
  17. from types import MappingProxyType
  18. from .dt import datetime_to_local_str, utcnow
  19. RE_SANITIZE_FILENAME = re.compile(r'(~|\.\.|/|\\)')
  20. RE_SANITIZE_PATH = re.compile(r'(~|\.(\.)+)')
  21. RE_SLUGIFY = re.compile(r'[^a-z0-9_]+')
  22. def sanitize_filename(filename):
  23. """ Sanitizes a filename by removing .. / and \\. """
  24. return RE_SANITIZE_FILENAME.sub("", filename)
  25. def sanitize_path(path):
  26. """ Sanitizes a path by removing ~ and .. """
  27. return RE_SANITIZE_PATH.sub("", path)
  28. def slugify(text):
  29. """ Slugifies a given text. """
  30. text = text.lower().replace(" ", "_")
  31. return RE_SLUGIFY.sub("", text)
  32. def repr_helper(inp):
  33. """ Helps creating a more readable string representation of objects. """
  34. if isinstance(inp, (dict, MappingProxyType)):
  35. return ", ".join(
  36. repr_helper(key)+"="+repr_helper(item) for key, item
  37. in inp.items())
  38. elif isinstance(inp, datetime):
  39. return datetime_to_local_str(inp)
  40. else:
  41. return str(inp)
  42. def convert(value, to_type, default=None):
  43. """ Converts value to to_type, returns default if fails. """
  44. try:
  45. return default if value is None else to_type(value)
  46. except (ValueError, TypeError):
  47. # If value could not be converted
  48. return default
  49. def ensure_unique_string(preferred_string, current_strings):
  50. """ Returns a string that is not present in current_strings.
  51. If preferred string exists will append _2, _3, .. """
  52. test_string = preferred_string
  53. current_strings = set(current_strings)
  54. tries = 1
  55. while test_string in current_strings:
  56. tries += 1
  57. test_string = "{}_{}".format(preferred_string, tries)
  58. return test_string
  59. # Taken from: http://stackoverflow.com/a/11735897
  60. def get_local_ip():
  61. """ Tries to determine the local IP address of the machine. """
  62. try:
  63. sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
  64. # Use Google Public DNS server to determine own IP
  65. sock.connect(('8.8.8.8', 80))
  66. return sock.getsockname()[0]
  67. except socket.error:
  68. return socket.gethostbyname(socket.gethostname())
  69. finally:
  70. sock.close()
  71. # Taken from http://stackoverflow.com/a/23728630
  72. def get_random_string(length=10):
  73. """ Returns a random string with letters and digits. """
  74. generator = random.SystemRandom()
  75. source_chars = string.ascii_letters + string.digits
  76. return ''.join(generator.choice(source_chars) for _ in range(length))
  77. class OrderedEnum(enum.Enum):
  78. """ Taken from Python 3.4.0 docs. """
  79. # pylint: disable=no-init, too-few-public-methods
  80. def __ge__(self, other):
  81. if self.__class__ is other.__class__:
  82. return self.value >= other.value
  83. return NotImplemented
  84. def __gt__(self, other):
  85. if self.__class__ is other.__class__:
  86. return self.value > other.value
  87. return NotImplemented
  88. def __le__(self, other):
  89. if self.__class__ is other.__class__:
  90. return self.value <= other.value
  91. return NotImplemented
  92. def __lt__(self, other):
  93. if self.__class__ is other.__class__:
  94. return self.value < other.value
  95. return NotImplemented
  96. class OrderedSet(collections.MutableSet):
  97. """ Ordered set taken from http://code.activestate.com/recipes/576694/ """
  98. def __init__(self, iterable=None):
  99. self.end = end = []
  100. end += [None, end, end] # sentinel node for doubly linked list
  101. self.map = {} # key --> [key, prev, next]
  102. if iterable is not None:
  103. self |= iterable
  104. def __len__(self):
  105. return len(self.map)
  106. def __contains__(self, key):
  107. return key in self.map
  108. def add(self, key):
  109. """ Add an element to the end of the set. """
  110. if key not in self.map:
  111. end = self.end
  112. curr = end[1]
  113. curr[2] = end[1] = self.map[key] = [key, curr, end]
  114. def promote(self, key):
  115. """ Promote element to beginning of the set, add if not there. """
  116. if key in self.map:
  117. self.discard(key)
  118. begin = self.end[2]
  119. curr = begin[1]
  120. curr[2] = begin[1] = self.map[key] = [key, curr, begin]
  121. def discard(self, key):
  122. """ Discard an element from the set. """
  123. if key in self.map:
  124. key, prev_item, next_item = self.map.pop(key)
  125. prev_item[2] = next_item
  126. next_item[1] = prev_item
  127. def __iter__(self):
  128. end = self.end
  129. curr = end[2]
  130. while curr is not end:
  131. yield curr[0]
  132. curr = curr[2]
  133. def __reversed__(self):
  134. end = self.end
  135. curr = end[1]
  136. while curr is not end:
  137. yield curr[0]
  138. curr = curr[1]
  139. def pop(self, last=True): # pylint: disable=arguments-differ
  140. """
  141. Pops element of the end of the set.
  142. Set last=False to pop from the beginning.
  143. """
  144. if not self:
  145. raise KeyError('set is empty')
  146. key = self.end[1][0] if last else self.end[2][0]
  147. self.discard(key)
  148. return key
  149. def update(self, *args):
  150. """ Add elements from args to the set. """
  151. for item in chain(*args):
  152. self.add(item)
  153. def __repr__(self):
  154. if not self:
  155. return '%s()' % (self.__class__.__name__,)
  156. return '%s(%r)' % (self.__class__.__name__, list(self))
  157. def __eq__(self, other):
  158. if isinstance(other, OrderedSet):
  159. return len(self) == len(other) and list(self) == list(other)
  160. return set(self) == set(other)
  161. class Throttle(object):
  162. """
  163. A method decorator to add a cooldown to a method to prevent it from being
  164. called more then 1 time within the timedelta interval `min_time` after it
  165. returned its result.
  166. Calling a method a second time during the interval will return None.
  167. Pass keyword argument `no_throttle=True` to the wrapped method to make
  168. the call not throttled.
  169. Decorator takes in an optional second timedelta interval to throttle the
  170. 'no_throttle' calls.
  171. Adds a datetime attribute `last_call` to the method.
  172. """
  173. # pylint: disable=too-few-public-methods
  174. def __init__(self, min_time, limit_no_throttle=None):
  175. self.min_time = min_time
  176. self.limit_no_throttle = limit_no_throttle
  177. def __call__(self, method):
  178. if self.limit_no_throttle is not None:
  179. method = Throttle(self.limit_no_throttle)(method)
  180. # Different methods that can be passed in:
  181. # - a function
  182. # - an unbound function on a class
  183. # - a method (bound function on a class)
  184. # We want to be able to differentiate between function and unbound
  185. # methods (which are considered functions).
  186. # All methods have the classname in their qualname seperated by a '.'
  187. # Functions have a '.' in their qualname if defined inline, but will
  188. # be prefixed by '.<locals>.' so we strip that out.
  189. is_func = (not hasattr(method, '__self__') and
  190. '.' not in method.__qualname__.split('.<locals>.')[-1])
  191. @wraps(method)
  192. def wrapper(*args, **kwargs):
  193. """
  194. Wrapper that allows wrapped to be called only once per min_time.
  195. If we cannot acquire the lock, it is running so return None.
  196. """
  197. # pylint: disable=protected-access
  198. if hasattr(method, '__self__'):
  199. host = method.__self__
  200. elif is_func:
  201. host = wrapper
  202. else:
  203. host = args[0] if args else wrapper
  204. if not hasattr(host, '_throttle'):
  205. host._throttle = {}
  206. if id(self) not in host._throttle:
  207. host._throttle[id(self)] = [threading.Lock(), None]
  208. throttle = host._throttle[id(self)]
  209. if not throttle[0].acquire(False):
  210. return None
  211. # Check if method is never called or no_throttle is given
  212. force = not throttle[1] or kwargs.pop('no_throttle', False)
  213. try:
  214. if force or utcnow() - throttle[1] > self.min_time:
  215. result = method(*args, **kwargs)
  216. throttle[1] = utcnow()
  217. return result
  218. else:
  219. return None
  220. finally:
  221. throttle[0].release()
  222. return wrapper
  223. class ThreadPool(object):
  224. """A priority queue-based thread pool."""
  225. # pylint: disable=too-many-instance-attributes
  226. def __init__(self, job_handler, worker_count=0, busy_callback=None):
  227. """
  228. job_handler: method to be called from worker thread to handle job
  229. worker_count: number of threads to run that handle jobs
  230. busy_callback: method to be called when queue gets too big.
  231. Parameters: worker_count, list of current_jobs,
  232. pending_jobs_count
  233. """
  234. self._job_handler = job_handler
  235. self._busy_callback = busy_callback
  236. self.worker_count = 0
  237. self.busy_warning_limit = 0
  238. self._work_queue = queue.PriorityQueue()
  239. self.current_jobs = []
  240. self._lock = threading.RLock()
  241. self._quit_task = object()
  242. self.running = True
  243. for _ in range(worker_count):
  244. self.add_worker()
  245. def add_worker(self):
  246. """Add worker to the thread pool and reset warning limit."""
  247. with self._lock:
  248. if not self.running:
  249. raise RuntimeError("ThreadPool not running")
  250. worker = threading.Thread(target=self._worker)
  251. worker.daemon = True
  252. worker.start()
  253. self.worker_count += 1
  254. self.busy_warning_limit = self.worker_count * 3
  255. def remove_worker(self):
  256. """Remove worker from the thread pool and reset warning limit."""
  257. with self._lock:
  258. if not self.running:
  259. raise RuntimeError("ThreadPool not running")
  260. self._work_queue.put(PriorityQueueItem(0, self._quit_task))
  261. self.worker_count -= 1
  262. self.busy_warning_limit = self.worker_count * 3
  263. def add_job(self, priority, job):
  264. """ Add a job to the queue. """
  265. with self._lock:
  266. if not self.running:
  267. raise RuntimeError("ThreadPool not running")
  268. self._work_queue.put(PriorityQueueItem(priority, job))
  269. # check if our queue is getting too big
  270. if self._work_queue.qsize() > self.busy_warning_limit \
  271. and self._busy_callback is not None:
  272. # Increase limit we will issue next warning
  273. self.busy_warning_limit *= 2
  274. self._busy_callback(
  275. self.worker_count, self.current_jobs,
  276. self._work_queue.qsize())
  277. def block_till_done(self):
  278. """Block till current work is done."""
  279. self._work_queue.join()
  280. # import traceback
  281. # traceback.print_stack()
  282. def stop(self):
  283. """Finish all the jobs and stops all the threads."""
  284. self.block_till_done()
  285. with self._lock:
  286. if not self.running:
  287. return
  288. # Tell the workers to quit
  289. for _ in range(self.worker_count):
  290. self.remove_worker()
  291. self.running = False
  292. # Wait till all workers have quit
  293. self.block_till_done()
  294. def _worker(self):
  295. """Handle jobs for the thread pool."""
  296. while True:
  297. # Get new item from work_queue
  298. job = self._work_queue.get().item
  299. if job == self._quit_task:
  300. self._work_queue.task_done()
  301. return
  302. # Add to current running jobs
  303. job_log = (utcnow(), job)
  304. self.current_jobs.append(job_log)
  305. # Do the job
  306. self._job_handler(job)
  307. # Remove from current running job
  308. self.current_jobs.remove(job_log)
  309. # Tell work_queue the task is done
  310. self._work_queue.task_done()
  311. class PriorityQueueItem(object):
  312. """ Holds a priority and a value. Used within PriorityQueue. """
  313. # pylint: disable=too-few-public-methods
  314. def __init__(self, priority, item):
  315. self.priority = priority
  316. self.item = item
  317. def __lt__(self, other):
  318. return self.priority < other.priority