A Python robot that edits Wikipedia and interacts with people over IRC https://en.wikipedia.org/wiki/User:EarwigBot

270 lines
11 KiB

  1. #! /usr/bin/env python
  2. # -*- coding: utf-8 -*-
  3. #
  4. # Copyright (C) 2009-2014 Ben Kurtovic <ben.kurtovic@gmail.com>
  5. #
  6. # Permission is hereby granted, free of charge, to any person obtaining a copy
  7. # of this software and associated documentation files (the "Software"), to deal
  8. # in the Software without restriction, including without limitation the rights
  9. # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
  10. # copies of the Software, and to permit persons to whom the Software is
  11. # furnished to do so, subject to the following conditions:
  12. #
  13. # The above copyright notice and this permission notice shall be included in
  14. # all copies or substantial portions of the Software.
  15. #
  16. # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
  17. # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
  18. # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
  19. # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
  20. # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
  21. # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
  22. # SOFTWARE.
  23. import imp
  24. from os import listdir, path
  25. from re import sub
  26. from threading import RLock, Thread
  27. from time import gmtime, strftime
  28. from earwigbot.commands import Command
  29. from earwigbot.tasks import Task
  30. __all__ = ["CommandManager", "TaskManager"]
  31. class _ResourceManager(object):
  32. """
  33. **EarwigBot: Resource Manager**
  34. Resources are essentially objects dynamically loaded by the bot, both
  35. packaged with it (built-in resources) and created by users (plugins, aka
  36. custom resources). Currently, the only two types of resources are IRC
  37. commands and bot tasks. These are both loaded from two locations: the
  38. :py:mod:`earwigbot.commands` and :py:mod:`earwigbot.tasks packages`, and
  39. the :file:`commands/` and :file:`tasks/` directories within the bot's
  40. working directory.
  41. This class handles the low-level tasks of (re)loading resources via
  42. :py:meth:`load`, retrieving specific resources via :py:meth:`get`, and
  43. iterating over all resources via :py:meth:`__iter__`.
  44. """
  45. def __init__(self, bot, name, base):
  46. self.bot = bot
  47. self.logger = bot.logger.getChild(name)
  48. self._resources = {}
  49. self._resource_name = name # e.g. "commands" or "tasks"
  50. self._resource_base = base # e.g. Command or Task
  51. self._resource_access_lock = RLock()
  52. def __repr__(self):
  53. """Return the canonical string representation of the manager."""
  54. res = "{0}(bot={1!r}, name={2!r}, base={3!r})"
  55. return res.format(self.__class__.__name__, self.bot,
  56. self._resource_name, self._resource_base)
  57. def __str__(self):
  58. """Return a nice string representation of the manager."""
  59. return "<{0} of {1}>".format(self.__class__.__name__, self.bot)
  60. def __iter__(self):
  61. with self.lock:
  62. for resource in self._resources.itervalues():
  63. yield resource
  64. def _load_resource(self, name, path, klass):
  65. """Instantiate a resource class and add it to the dictionary."""
  66. res_type = self._resource_name[:-1] # e.g. "command" or "task"
  67. if hasattr(klass, "name"):
  68. res_config = getattr(self.bot.config, self._resource_name)
  69. if getattr(klass, "name") in res_config.get("disable", []):
  70. log = "Skipping disabled {0} {1}"
  71. self.logger.debug(log.format(res_type, getattr(klass, "name")))
  72. return
  73. try:
  74. resource = klass(self.bot) # Create instance of resource
  75. except Exception:
  76. e = "Error instantiating {0} class in '{1}' (from {2})"
  77. self.logger.exception(e.format(res_type, name, path))
  78. else:
  79. self._resources[resource.name] = resource
  80. self.logger.debug("Loaded {0} {1}".format(res_type, resource.name))
  81. def _load_module(self, name, path):
  82. """Load a specific resource from a module, identified by name and path.
  83. We'll first try to import it using imp magic, and if that works, make
  84. instances of any classes inside that are subclasses of the base
  85. (:py:attr:`self._resource_base <_resource_base>`), add them to the
  86. resources dictionary with :py:meth:`self._load_resource()
  87. <_load_resource>`, and finally log the addition. Any problems along
  88. the way will either be ignored or logged.
  89. """
  90. f, path, desc = imp.find_module(name, [path])
  91. try:
  92. module = imp.load_module(name, f, path, desc)
  93. except Exception:
  94. e = "Couldn't load module '{0}' (from {1})"
  95. self.logger.exception(e.format(name, path))
  96. return
  97. finally:
  98. f.close()
  99. for obj in vars(module).values():
  100. if type(obj) is type:
  101. isresource = issubclass(obj, self._resource_base)
  102. if isresource and not obj is self._resource_base:
  103. self._load_resource(name, path, obj)
  104. def _load_directory(self, dir):
  105. """Load all valid resources in a given directory."""
  106. self.logger.debug("Loading directory {0}".format(dir))
  107. res_config = getattr(self.bot.config, self._resource_name)
  108. disabled = res_config.get("disable", [])
  109. processed = []
  110. for name in listdir(dir):
  111. if not name.endswith(".py") and not name.endswith(".pyc"):
  112. continue
  113. if name.startswith("_") or name.startswith("."):
  114. continue
  115. modname = sub("\.pyc?$", "", name) # Remove extension
  116. if modname in disabled:
  117. log = "Skipping disabled module {0}".format(modname)
  118. self.logger.debug(log)
  119. continue
  120. if modname not in processed:
  121. self._load_module(modname, dir)
  122. processed.append(modname)
  123. @property
  124. def lock(self):
  125. """The resource access/modify lock."""
  126. return self._resource_access_lock
  127. def load(self):
  128. """Load (or reload) all valid resources into :py:attr:`_resources`."""
  129. name = self._resource_name # e.g. "commands" or "tasks"
  130. with self.lock:
  131. self._resources.clear()
  132. builtin_dir = path.join(path.dirname(__file__), name)
  133. plugins_dir = path.join(self.bot.config.root_dir, name)
  134. if getattr(self.bot.config, name).get("disable") is True:
  135. log = "Skipping disabled builtins directory: {0}"
  136. self.logger.debug(log.format(builtin_dir))
  137. else:
  138. self._load_directory(builtin_dir) # Built-in resources
  139. if path.exists(plugins_dir) and path.isdir(plugins_dir):
  140. self._load_directory(plugins_dir) # Custom resources, plugins
  141. else:
  142. log = "Skipping nonexistent plugins directory: {0}"
  143. self.logger.debug(log.format(plugins_dir))
  144. if self._resources:
  145. msg = "Loaded {0} {1}: {2}"
  146. resources = ", ".join(self._resources.keys())
  147. self.logger.info(msg.format(len(self._resources), name, resources))
  148. else:
  149. self.logger.info("Loaded 0 {0}".format(name))
  150. def get(self, key):
  151. """Return the class instance associated with a certain resource.
  152. Will raise :py:exc:`KeyError` if the resource (a command or task) is
  153. not found.
  154. """
  155. with self.lock:
  156. return self._resources[key]
  157. class CommandManager(_ResourceManager):
  158. """
  159. Manages (i.e., loads, reloads, and calls) IRC commands.
  160. """
  161. def __init__(self, bot):
  162. super(CommandManager, self).__init__(bot, "commands", Command)
  163. def _wrap_check(self, command, data):
  164. """Check whether a command should be called, catching errors."""
  165. try:
  166. return command.check(data)
  167. except Exception:
  168. e = "Error checking command '{0}' with data: {1}:"
  169. self.logger.exception(e.format(command.name, data))
  170. def _wrap_process(self, command, data):
  171. """process() the message, catching and reporting any errors."""
  172. try:
  173. command.process(data)
  174. except Exception:
  175. e = "Error executing command '{0}':"
  176. self.logger.exception(e.format(command.name))
  177. def call(self, hook, data):
  178. """Respond to a hook type and a :py:class:`Data` object."""
  179. for command in self:
  180. if hook in command.hooks and self._wrap_check(command, data):
  181. thread = Thread(target=self._wrap_process,
  182. args=(command, data))
  183. start_time = strftime("%b %d %H:%M:%S")
  184. thread.name = "irc:{0} ({1})".format(command.name, start_time)
  185. thread.daemon = True
  186. thread.start()
  187. return
  188. class TaskManager(_ResourceManager):
  189. """
  190. Manages (i.e., loads, reloads, schedules, and runs) wiki bot tasks.
  191. """
  192. def __init__(self, bot):
  193. super(TaskManager, self).__init__(bot, "tasks", Task)
  194. def _wrapper(self, task, **kwargs):
  195. """Wrapper for task classes: run the task and catch any errors."""
  196. try:
  197. task.run(**kwargs)
  198. except Exception:
  199. msg = "Task '{0}' raised an exception and had to stop:"
  200. self.logger.exception(msg.format(task.name))
  201. else:
  202. msg = "Task '{0}' finished successfully"
  203. self.logger.info(msg.format(task.name))
  204. def start(self, task_name, **kwargs):
  205. """Start a given task in a new daemon thread, and return the thread.
  206. kwargs are passed to :py:meth:`task.run() <earwigbot.tasks.Task.run>`.
  207. If the task is not found, ``None`` will be returned and an error will
  208. be logged.
  209. """
  210. msg = "Starting task '{0}' in a new thread"
  211. self.logger.info(msg.format(task_name))
  212. try:
  213. task = self.get(task_name)
  214. except KeyError:
  215. e = "Couldn't find task '{0}'"
  216. self.logger.error(e.format(task_name))
  217. return
  218. task_thread = Thread(target=self._wrapper, args=(task,), kwargs=kwargs)
  219. start_time = strftime("%b %d %H:%M:%S")
  220. task_thread.name = "{0} ({1})".format(task_name, start_time)
  221. task_thread.daemon = True
  222. task_thread.start()
  223. return task_thread
  224. def schedule(self, now=None):
  225. """Start all tasks that are supposed to be run at a given time."""
  226. if not now:
  227. now = gmtime()
  228. # Get list of tasks to run this turn:
  229. tasks = self.bot.config.schedule(now.tm_min, now.tm_hour, now.tm_mday,
  230. now.tm_mon, now.tm_wday)
  231. for task in tasks:
  232. if isinstance(task, list): # They've specified kwargs,
  233. self.start(task[0], **task[1]) # so pass those to start
  234. else: # Otherwise, just pass task_name
  235. self.start(task)