utils.py 14 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416
  1. """
  2. Internal subroutines for e.g. aborting execution with an error message,
  3. or performing indenting on multiline output.
  4. """
  5. import os
  6. import sys
  7. import textwrap
  8. from traceback import format_exc
  9. def _encode(msg, stream):
  10. if isinstance(msg, unicode) and hasattr(stream, 'encoding') and not stream.encoding is None:
  11. return msg.encode(stream.encoding)
  12. else:
  13. return str(msg)
  14. def isatty(stream):
  15. """Check if a stream is a tty.
  16. Not all file-like objects implement the `isatty` method.
  17. """
  18. fn = getattr(stream, 'isatty', None)
  19. if fn is None:
  20. return False
  21. return fn()
  22. def abort(msg):
  23. """
  24. Abort execution, print ``msg`` to stderr and exit with error status (1.)
  25. This function currently makes use of `SystemExit`_ in a manner that is
  26. similar to `sys.exit`_ (but which skips the automatic printing to stderr,
  27. allowing us to more tightly control it via settings).
  28. Therefore, it's possible to detect and recover from inner calls to `abort`
  29. by using ``except SystemExit`` or similar.
  30. .. _sys.exit: http://docs.python.org/library/sys.html#sys.exit
  31. .. _SystemExit: http://docs.python.org/library/exceptions.html#exceptions.SystemExit
  32. """
  33. from fabric.state import output, env
  34. if not env.colorize_errors:
  35. red = lambda x: x
  36. else:
  37. from colors import red
  38. if output.aborts:
  39. sys.stderr.write(red("\nFatal error: %s\n" % _encode(msg, sys.stderr)))
  40. sys.stderr.write(red("\nAborting.\n"))
  41. if env.abort_exception:
  42. raise env.abort_exception(msg)
  43. else:
  44. # See issue #1318 for details on the below; it lets us construct a
  45. # valid, useful SystemExit while sidestepping the automatic stderr
  46. # print (which would otherwise duplicate with the above in a
  47. # non-controllable fashion).
  48. e = SystemExit(1)
  49. e.message = msg
  50. raise e
  51. def warn(msg):
  52. """
  53. Print warning message, but do not abort execution.
  54. This function honors Fabric's :doc:`output controls
  55. <../../usage/output_controls>` and will print the given ``msg`` to stderr,
  56. provided that the ``warnings`` output level (which is active by default) is
  57. turned on.
  58. """
  59. from fabric.state import output, env
  60. if not env.colorize_errors:
  61. magenta = lambda x: x
  62. else:
  63. from colors import magenta
  64. if output.warnings:
  65. msg = _encode(msg, sys.stderr)
  66. sys.stderr.write(magenta("\nWarning: %s\n\n" % msg))
  67. def indent(text, spaces=4, strip=False):
  68. """
  69. Return ``text`` indented by the given number of spaces.
  70. If text is not a string, it is assumed to be a list of lines and will be
  71. joined by ``\\n`` prior to indenting.
  72. When ``strip`` is ``True``, a minimum amount of whitespace is removed from
  73. the left-hand side of the given string (so that relative indents are
  74. preserved, but otherwise things are left-stripped). This allows you to
  75. effectively "normalize" any previous indentation for some inputs.
  76. """
  77. # Normalize list of strings into a string for dedenting. "list" here means
  78. # "not a string" meaning "doesn't have splitlines". Meh.
  79. if not hasattr(text, 'splitlines'):
  80. text = '\n'.join(text)
  81. # Dedent if requested
  82. if strip:
  83. text = textwrap.dedent(text)
  84. prefix = ' ' * spaces
  85. output = '\n'.join(prefix + line for line in text.splitlines())
  86. # Strip out empty lines before/aft
  87. output = output.strip()
  88. # Reintroduce first indent (which just got stripped out)
  89. output = prefix + output
  90. return output
  91. def puts(text, show_prefix=None, end="\n", flush=False):
  92. """
  93. An alias for ``print`` whose output is managed by Fabric's output controls.
  94. In other words, this function simply prints to ``sys.stdout``, but will
  95. hide its output if the ``user`` :doc:`output level
  96. </usage/output_controls>` is set to ``False``.
  97. If ``show_prefix=False``, `puts` will omit the leading ``[hostname]``
  98. which it tacks on by default. (It will also omit this prefix if
  99. ``env.host_string`` is empty.)
  100. Newlines may be disabled by setting ``end`` to the empty string (``''``).
  101. (This intentionally mirrors Python 3's ``print`` syntax.)
  102. You may force output flushing (e.g. to bypass output buffering) by setting
  103. ``flush=True``.
  104. .. versionadded:: 0.9.2
  105. .. seealso:: `~fabric.utils.fastprint`
  106. """
  107. from fabric.state import output, env
  108. if show_prefix is None:
  109. show_prefix = env.output_prefix
  110. if output.user:
  111. prefix = ""
  112. if env.host_string and show_prefix:
  113. prefix = "[%s] " % env.host_string
  114. sys.stdout.write(prefix + _encode(text, sys.stdout) + end)
  115. if flush:
  116. sys.stdout.flush()
  117. def fastprint(text, show_prefix=False, end="", flush=True):
  118. """
  119. Print ``text`` immediately, without any prefix or line ending.
  120. This function is simply an alias of `~fabric.utils.puts` with different
  121. default argument values, such that the ``text`` is printed without any
  122. embellishment and immediately flushed.
  123. It is useful for any situation where you wish to print text which might
  124. otherwise get buffered by Python's output buffering (such as within a
  125. processor intensive ``for`` loop). Since such use cases typically also
  126. require a lack of line endings (such as printing a series of dots to
  127. signify progress) it also omits the traditional newline by default.
  128. .. note::
  129. Since `~fabric.utils.fastprint` calls `~fabric.utils.puts`, it is
  130. likewise subject to the ``user`` :doc:`output level
  131. </usage/output_controls>`.
  132. .. versionadded:: 0.9.2
  133. .. seealso:: `~fabric.utils.puts`
  134. """
  135. return puts(text=text, show_prefix=show_prefix, end=end, flush=flush)
  136. def handle_prompt_abort(prompt_for):
  137. import fabric.state
  138. reason = "Needed to prompt for %s (host: %s), but %%s" % (
  139. prompt_for, fabric.state.env.host_string
  140. )
  141. # Explicit "don't prompt me bro"
  142. if fabric.state.env.abort_on_prompts:
  143. abort(reason % "abort-on-prompts was set to True")
  144. # Implicit "parallel == stdin/prompts have ambiguous target"
  145. if fabric.state.env.parallel:
  146. abort(reason % "input would be ambiguous in parallel mode")
  147. class _AttributeDict(dict):
  148. """
  149. Dictionary subclass enabling attribute lookup/assignment of keys/values.
  150. For example::
  151. >>> m = _AttributeDict({'foo': 'bar'})
  152. >>> m.foo
  153. 'bar'
  154. >>> m.foo = 'not bar'
  155. >>> m['foo']
  156. 'not bar'
  157. ``_AttributeDict`` objects also provide ``.first()`` which acts like
  158. ``.get()`` but accepts multiple keys as arguments, and returns the value of
  159. the first hit, e.g.::
  160. >>> m = _AttributeDict({'foo': 'bar', 'biz': 'baz'})
  161. >>> m.first('wrong', 'incorrect', 'foo', 'biz')
  162. 'bar'
  163. """
  164. def __getattr__(self, key):
  165. try:
  166. return self[key]
  167. except KeyError:
  168. # to conform with __getattr__ spec
  169. raise AttributeError(key)
  170. def __setattr__(self, key, value):
  171. self[key] = value
  172. def first(self, *names):
  173. for name in names:
  174. value = self.get(name)
  175. if value:
  176. return value
  177. class _AliasDict(_AttributeDict):
  178. """
  179. `_AttributeDict` subclass that allows for "aliasing" of keys to other keys.
  180. Upon creation, takes an ``aliases`` mapping, which should map alias names
  181. to lists of key names. Aliases do not store their own value, but instead
  182. set (override) all mapped keys' values. For example, in the following
  183. `_AliasDict`, calling ``mydict['foo'] = True`` will set the values of
  184. ``mydict['bar']``, ``mydict['biz']`` and ``mydict['baz']`` all to True::
  185. mydict = _AliasDict(
  186. {'biz': True, 'baz': False},
  187. aliases={'foo': ['bar', 'biz', 'baz']}
  188. )
  189. Because it is possible for the aliased values to be in a heterogenous
  190. state, reading aliases is not supported -- only writing to them is allowed.
  191. This also means they will not show up in e.g. ``dict.keys()``.
  192. ..note::
  193. Aliases are recursive, so you may refer to an alias within the key list
  194. of another alias. Naturally, this means that you can end up with
  195. infinite loops if you're not careful.
  196. `_AliasDict` provides a special function, `expand_aliases`, which will take
  197. a list of keys as an argument and will return that list of keys with any
  198. aliases expanded. This function will **not** dedupe, so any aliases which
  199. overlap will result in duplicate keys in the resulting list.
  200. """
  201. def __init__(self, arg=None, aliases=None):
  202. init = super(_AliasDict, self).__init__
  203. if arg is not None:
  204. init(arg)
  205. else:
  206. init()
  207. # Can't use super() here because of _AttributeDict's setattr override
  208. dict.__setattr__(self, 'aliases', aliases)
  209. def __setitem__(self, key, value):
  210. # Attr test required to not blow up when deepcopy'd
  211. if hasattr(self, 'aliases') and key in self.aliases:
  212. for aliased in self.aliases[key]:
  213. self[aliased] = value
  214. else:
  215. return super(_AliasDict, self).__setitem__(key, value)
  216. def expand_aliases(self, keys):
  217. ret = []
  218. for key in keys:
  219. if key in self.aliases:
  220. ret.extend(self.expand_aliases(self.aliases[key]))
  221. else:
  222. ret.append(key)
  223. return ret
  224. def _pty_size():
  225. """
  226. Obtain (rows, cols) tuple for sizing a pty on the remote end.
  227. Defaults to 80x24 (which is also the 'ssh' lib's default) but will detect
  228. local (stdout-based) terminal window size on non-Windows platforms.
  229. """
  230. from fabric.state import win32
  231. if not win32:
  232. import fcntl
  233. import termios
  234. import struct
  235. default_rows, default_cols = 24, 80
  236. rows, cols = default_rows, default_cols
  237. if not win32 and isatty(sys.stdout):
  238. # We want two short unsigned integers (rows, cols)
  239. fmt = 'HH'
  240. # Create an empty (zeroed) buffer for ioctl to map onto. Yay for C!
  241. buffer = struct.pack(fmt, 0, 0)
  242. # Call TIOCGWINSZ to get window size of stdout, returns our filled
  243. # buffer
  244. try:
  245. result = fcntl.ioctl(sys.stdout.fileno(), termios.TIOCGWINSZ,
  246. buffer)
  247. # Unpack buffer back into Python data types
  248. rows, cols = struct.unpack(fmt, result)
  249. # Fall back to defaults if TIOCGWINSZ returns unreasonable values
  250. if rows == 0:
  251. rows = default_rows
  252. if cols == 0:
  253. cols = default_cols
  254. # Deal with e.g. sys.stdout being monkeypatched, such as in testing.
  255. # Or termios not having a TIOCGWINSZ.
  256. except AttributeError:
  257. pass
  258. return rows, cols
  259. def error(message, func=None, exception=None, stdout=None, stderr=None):
  260. """
  261. Call ``func`` with given error ``message``.
  262. If ``func`` is None (the default), the value of ``env.warn_only``
  263. determines whether to call ``abort`` or ``warn``.
  264. If ``exception`` is given, it is inspected to get a string message, which
  265. is printed alongside the user-generated ``message``.
  266. If ``stdout`` and/or ``stderr`` are given, they are assumed to be strings
  267. to be printed.
  268. """
  269. import fabric.state
  270. if func is None:
  271. func = fabric.state.env.warn_only and warn or abort
  272. # If exception printing is on, append a traceback to the message
  273. if fabric.state.output.exceptions or fabric.state.output.debug:
  274. exception_message = format_exc()
  275. if exception_message:
  276. message += "\n\n" + exception_message
  277. # Otherwise, if we were given an exception, append its contents.
  278. elif exception is not None:
  279. # Figure out how to get a string out of the exception; EnvironmentError
  280. # subclasses, for example, "are" integers and .strerror is the string.
  281. # Others "are" strings themselves. May have to expand this further for
  282. # other error types.
  283. if hasattr(exception, 'strerror') and exception.strerror is not None:
  284. underlying = exception.strerror
  285. else:
  286. underlying = exception
  287. message += "\n\nUnderlying exception:\n" + indent(str(underlying))
  288. if func is abort:
  289. if stdout and not fabric.state.output.stdout:
  290. message += _format_error_output("Standard output", stdout)
  291. if stderr and not fabric.state.output.stderr:
  292. message += _format_error_output("Standard error", stderr)
  293. return func(message)
  294. def _format_error_output(header, body):
  295. term_width = _pty_size()[1]
  296. header_side_length = (term_width - (len(header) + 2)) / 2
  297. mark = "="
  298. side = mark * header_side_length
  299. return "\n\n%s %s %s\n\n%s\n\n%s" % (
  300. side, header, side, body, mark * term_width
  301. )
  302. # TODO: replace with collections.deque(maxlen=xxx) in Python 2.6
  303. class RingBuffer(list):
  304. def __init__(self, value, maxlen):
  305. # Because it's annoying typing this multiple times.
  306. self._super = super(RingBuffer, self)
  307. # Python 2.6 deque compatible option name!
  308. self._maxlen = maxlen
  309. return self._super.__init__(value)
  310. def _trim(self):
  311. if self._maxlen is None:
  312. return
  313. overage = max(len(self) - self._maxlen, 0)
  314. del self[0:overage]
  315. def append(self, value):
  316. self._super.append(value)
  317. self._trim()
  318. def extend(self, values):
  319. self._super.extend(values)
  320. self._trim()
  321. def __iadd__(self, other):
  322. self.extend(other)
  323. return self
  324. # Paranoia from here on out.
  325. def insert(self, index, value):
  326. raise ValueError("Can't insert into the middle of a ring buffer!")
  327. def __setslice__(self, i, j, sequence):
  328. raise ValueError("Can't set a slice of a ring buffer!")
  329. def __setitem__(self, key, value):
  330. if isinstance(key, slice):
  331. raise ValueError("Can't set a slice of a ring buffer!")
  332. else:
  333. return self._super.__setitem__(key, value)
  334. def apply_lcwd(path, env):
  335. # Apply CWD if a relative path
  336. if not os.path.isabs(path) and env.lcwd:
  337. path = os.path.join(env.lcwd, path)
  338. return path