decorators.py 7.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222
  1. """
  2. Convenience decorators for use in fabfiles.
  3. """
  4. from __future__ import with_statement
  5. import types
  6. from functools import wraps
  7. try:
  8. from Crypto import Random
  9. except ImportError:
  10. Random = None
  11. from fabric import tasks
  12. from .context_managers import settings
  13. def task(*args, **kwargs):
  14. """
  15. Decorator declaring the wrapped function to be a new-style task.
  16. May be invoked as a simple, argument-less decorator (i.e. ``@task``) or
  17. with arguments customizing its behavior (e.g. ``@task(alias='myalias')``).
  18. Please see the :ref:`new-style task <task-decorator>` documentation for
  19. details on how to use this decorator.
  20. .. versionchanged:: 1.2
  21. Added the ``alias``, ``aliases``, ``task_class`` and ``default``
  22. keyword arguments. See :ref:`task-decorator-arguments` for details.
  23. .. versionchanged:: 1.5
  24. Added the ``name`` keyword argument.
  25. .. seealso:: `~fabric.docs.unwrap_tasks`, `~fabric.tasks.WrappedCallableTask`
  26. """
  27. invoked = bool(not args or kwargs)
  28. task_class = kwargs.pop("task_class", tasks.WrappedCallableTask)
  29. if not invoked:
  30. func, args = args[0], ()
  31. def wrapper(func):
  32. return task_class(func, *args, **kwargs)
  33. return wrapper if invoked else wrapper(func)
  34. def _wrap_as_new(original, new):
  35. if isinstance(original, tasks.Task):
  36. return tasks.WrappedCallableTask(new)
  37. return new
  38. def _list_annotating_decorator(attribute, *values):
  39. def attach_list(func):
  40. @wraps(func)
  41. def inner_decorator(*args, **kwargs):
  42. return func(*args, **kwargs)
  43. _values = values
  44. # Allow for single iterable argument as well as *args
  45. if len(_values) == 1 and not isinstance(_values[0], basestring):
  46. _values = _values[0]
  47. setattr(inner_decorator, attribute, list(_values))
  48. # Don't replace @task new-style task objects with inner_decorator by
  49. # itself -- wrap in a new Task object first.
  50. inner_decorator = _wrap_as_new(func, inner_decorator)
  51. return inner_decorator
  52. return attach_list
  53. def hosts(*host_list):
  54. """
  55. Decorator defining which host or hosts to execute the wrapped function on.
  56. For example, the following will ensure that, barring an override on the
  57. command line, ``my_func`` will be run on ``host1``, ``host2`` and
  58. ``host3``, and with specific users on ``host1`` and ``host3``::
  59. @hosts('user1@host1', 'host2', 'user2@host3')
  60. def my_func():
  61. pass
  62. `~fabric.decorators.hosts` may be invoked with either an argument list
  63. (``@hosts('host1')``, ``@hosts('host1', 'host2')``) or a single, iterable
  64. argument (``@hosts(['host1', 'host2'])``).
  65. Note that this decorator actually just sets the function's ``.hosts``
  66. attribute, which is then read prior to executing the function.
  67. .. versionchanged:: 0.9.2
  68. Allow a single, iterable argument (``@hosts(iterable)``) to be used
  69. instead of requiring ``@hosts(*iterable)``.
  70. """
  71. return _list_annotating_decorator('hosts', *host_list)
  72. def roles(*role_list):
  73. """
  74. Decorator defining a list of role names, used to look up host lists.
  75. A role is simply defined as a key in `env` whose value is a list of one or
  76. more host connection strings. For example, the following will ensure that,
  77. barring an override on the command line, ``my_func`` will be executed
  78. against the hosts listed in the ``webserver`` and ``dbserver`` roles::
  79. env.roledefs.update({
  80. 'webserver': ['www1', 'www2'],
  81. 'dbserver': ['db1']
  82. })
  83. @roles('webserver', 'dbserver')
  84. def my_func():
  85. pass
  86. As with `~fabric.decorators.hosts`, `~fabric.decorators.roles` may be
  87. invoked with either an argument list or a single, iterable argument.
  88. Similarly, this decorator uses the same mechanism as
  89. `~fabric.decorators.hosts` and simply sets ``<function>.roles``.
  90. .. versionchanged:: 0.9.2
  91. Allow a single, iterable argument to be used (same as
  92. `~fabric.decorators.hosts`).
  93. """
  94. return _list_annotating_decorator('roles', *role_list)
  95. def runs_once(func):
  96. """
  97. Decorator preventing wrapped function from running more than once.
  98. By keeping internal state, this decorator allows you to mark a function
  99. such that it will only run once per Python interpreter session, which in
  100. typical use means "once per invocation of the ``fab`` program".
  101. Any function wrapped with this decorator will silently fail to execute the
  102. 2nd, 3rd, ..., Nth time it is called, and will return the value of the
  103. original run.
  104. .. note:: ``runs_once`` does not work with parallel task execution.
  105. """
  106. @wraps(func)
  107. def decorated(*args, **kwargs):
  108. if not hasattr(decorated, 'return_value'):
  109. decorated.return_value = func(*args, **kwargs)
  110. return decorated.return_value
  111. decorated = _wrap_as_new(func, decorated)
  112. # Mark as serial (disables parallelism) and return
  113. return serial(decorated)
  114. def serial(func):
  115. """
  116. Forces the wrapped function to always run sequentially, never in parallel.
  117. This decorator takes precedence over the global value of :ref:`env.parallel
  118. <env-parallel>`. However, if a task is decorated with both
  119. `~fabric.decorators.serial` *and* `~fabric.decorators.parallel`,
  120. `~fabric.decorators.parallel` wins.
  121. .. versionadded:: 1.3
  122. """
  123. if not getattr(func, 'parallel', False):
  124. func.serial = True
  125. return _wrap_as_new(func, func)
  126. def parallel(pool_size=None):
  127. """
  128. Forces the wrapped function to run in parallel, instead of sequentially.
  129. This decorator takes precedence over the global value of :ref:`env.parallel
  130. <env-parallel>`. It also takes precedence over `~fabric.decorators.serial`
  131. if a task is decorated with both.
  132. .. versionadded:: 1.3
  133. """
  134. called_without_args = type(pool_size) == types.FunctionType
  135. def real_decorator(func):
  136. @wraps(func)
  137. def inner(*args, **kwargs):
  138. # Required for ssh/PyCrypto to be happy in multiprocessing
  139. # (as far as we can tell, this is needed even with the extra such
  140. # calls in newer versions of paramiko.)
  141. if Random:
  142. Random.atfork()
  143. return func(*args, **kwargs)
  144. inner.parallel = True
  145. inner.serial = False
  146. inner.pool_size = None if called_without_args else pool_size
  147. return _wrap_as_new(func, inner)
  148. # Allow non-factory-style decorator use (@decorator vs @decorator())
  149. if called_without_args:
  150. return real_decorator(pool_size)
  151. return real_decorator
  152. def with_settings(*arg_settings, **kw_settings):
  153. """
  154. Decorator equivalent of ``fabric.context_managers.settings``.
  155. Allows you to wrap an entire function as if it was called inside a block
  156. with the ``settings`` context manager. This may be useful if you know you
  157. want a given setting applied to an entire function body, or wish to
  158. retrofit old code without indenting everything.
  159. For example, to turn aborts into warnings for an entire task function::
  160. @with_settings(warn_only=True)
  161. def foo():
  162. ...
  163. .. seealso:: `~fabric.context_managers.settings`
  164. .. versionadded:: 1.1
  165. """
  166. def outer(func):
  167. @wraps(func)
  168. def inner(*args, **kwargs):
  169. with settings(*arg_settings, **kw_settings):
  170. return func(*args, **kwargs)
  171. return _wrap_as_new(func, inner)
  172. return outer