A Python robot that edits Wikipedia and interacts with people over IRC https://en.wikipedia.org/wiki/User:EarwigBot
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

__init__.py 3.5 KiB

13 vuotta sitten
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118
  1. # -*- coding: utf-8 -*-
  2. """
  3. EarwigBot's Wiki Task Manager
  4. This package provides the wiki bot "tasks" EarwigBot runs. Here in __init__,
  5. you can find some functions used to load and run these tasks.
  6. """
  7. import logging
  8. import os
  9. import sys
  10. import threading
  11. import time
  12. from classes import BaseTask
  13. import config
  14. __all__ = ["load", "schedule", "start", "get", "get_all"]
  15. # Base directory when searching for tasks:
  16. base_dir = os.path.join(config.root_dir, "bot", "tasks")
  17. # Store loaded tasks as a dict where the key is the task name and the value is
  18. # an instance of the task class:
  19. _tasks = {}
  20. # Logger for this module:
  21. logger = logging.getLogger("commands")
  22. def _load_task(filename):
  23. """Try to load a specific task from a module, identified by file name."""
  24. global _tasks
  25. # Strip .py from the end of the filename and join with our package name:
  26. name = ".".join(("tasks", filename[:-3]))
  27. try:
  28. __import__(name)
  29. except:
  30. logger.exception("Couldn't load file {0}:".format(filename))
  31. return
  32. task = sys.modules[name].Task()
  33. task._setup_logger()
  34. if not isinstance(task, BaseTask):
  35. return
  36. _tasks[task.name] = task
  37. logger.debug("Added task {0}".format(task.name))
  38. def _wrapper(task, **kwargs):
  39. """Wrapper for task classes: run the task and catch any errors."""
  40. try:
  41. task.run(**kwargs)
  42. except:
  43. error = "Task '{0}' raised an exception and had to stop"
  44. logger.exception(error.format(task.name))
  45. else:
  46. logger.info("Task '{0}' finished without error".format(task.name))
  47. def load():
  48. """Load all valid tasks from bot/tasks/, into the _tasks variable."""
  49. files = os.listdir(base_dir)
  50. files.sort()
  51. for filename in files:
  52. if filename.startswith("_") or not filename.endswith(".py"):
  53. continue
  54. try:
  55. _load_task(filename)
  56. except AttributeError:
  57. pass # The file is doesn't contain a task, so just move on
  58. logger.info("Found {0} tasks: {1}".format(len(_tasks), ', '.join(_tasks.keys())))
  59. def schedule(now=time.gmtime()):
  60. """Start all tasks that are supposed to be run at a given time."""
  61. # Get list of tasks to run this turn:
  62. tasks = config.schedule(now.tm_min, now.tm_hour, now.tm_mday, now.tm_mon,
  63. now.tm_wday)
  64. for task in tasks:
  65. if isinstance(task, list): # they've specified kwargs
  66. start(task[0], **task[1]) # so pass those to start_task
  67. else: # otherwise, just pass task_name
  68. start(task)
  69. def start(task_name, **kwargs):
  70. """Start a given task in a new thread. Pass args to the task's run()
  71. function."""
  72. logger.info("Starting task '{0}' in a new thread".format(task_name))
  73. try:
  74. task = _tasks[task_name]
  75. except KeyError:
  76. error = "Couldn't find task '{0}': bot/tasks/{0}.py does not exist"
  77. logger.error(error.format(task_name))
  78. return
  79. task_thread = threading.Thread(target=lambda: _wrapper(task, **kwargs))
  80. start_time = time.strftime("%b %d %H:%M:%S")
  81. task_thread.name = "{0} ({1})".format(task_name, start_time)
  82. # Stop bot task threads automagically if the main bot stops:
  83. task_thread.daemon = True
  84. task_thread.start()
  85. def get(task_name):
  86. """Return the class instance associated with a certain task name.
  87. Will raise KeyError if the task is not found.
  88. """
  89. return _tasks[task_name]
  90. def get_all():
  91. """Return our dict of all loaded tasks."""
  92. return _tasks