| # -*- encoding: utf-8 -*- |
| # Copyright 2014 The LUCI Authors. All rights reserved. |
| # Use of this source code is governed under the Apache License, Version 2.0 |
| # that can be found in the LICENSE file. |
| |
| """Provides objects for reading and writing raw data to and from steps.""" |
| |
| from recipe_engine import recipe_api |
| from recipe_engine import util as recipe_util |
| |
| import codecs |
| import contextlib |
| import cStringIO |
| import os |
| import shutil |
| import sys |
| import tempfile |
| |
| |
| def _rmfile(p): # pragma: no cover |
| """Deletes a file, even a read-only one on Windows.""" |
| if sys.platform == 'win32': |
| try: |
| os.remove(p) |
| except OSError: |
| # Try to remove the read-only bit and remove again. |
| os.chmod(p, 0777) |
| os.remove(p) |
| else: |
| os.remove(p) |
| |
| |
| def _rmtree(d): # pragma: no cover |
| """Deletes a directory without throwing, even one with read-only files.""" |
| if not os.path.exists(d): |
| return |
| |
| if sys.platform == 'win32': |
| # Tested manually. |
| def unset_ro_and_remove_again(fn, p, excinfo): |
| """Removes file even if it has the READ_ONLY file attribute. |
| |
| On Windows, a file with the READ_ONLY file attribute cannot be deleted. |
| This is different than on POSIX where only the containing directory ACL |
| matters. |
| |
| shutil.rmtree() has trouble with this. Helps it a bit. |
| """ |
| # fn is one of islink, listdir, remove or rmdir. |
| if fn is os.remove: |
| # Try to remove the read-only bit. |
| os.chmod(p, 0777) |
| # And remove again. |
| os.remove(p) |
| return |
| # Reraise the original exception. |
| raise excinfo[0], excinfo[1], excinfo[2] |
| |
| # On Windows, some paths exceed MAX_PATH. Work around this by prepending |
| # the UNC magic prefix '\\?\' which allows the Windows API file calls to |
| # ignore the MAX_PATH limit. |
| shutil.rmtree(ur'\\?\%s' % (d,), onerror=unset_ro_and_remove_again) |
| else: |
| shutil.rmtree(d) |
| |
| |
| class InputDataPlaceholder(recipe_util.InputPlaceholder): |
| def __init__(self, data, suffix, name=None): |
| if not isinstance(data, str): # pragma: no cover |
| raise TypeError( |
| "Data passed to InputDataPlaceholder was %r, expected 'str'" |
| % (type(data).__name__)) |
| self.data = data |
| self.suffix = suffix |
| self._backing_file = None |
| super(InputDataPlaceholder, self).__init__(name=name) |
| |
| @property |
| def backing_file(self): |
| return self._backing_file |
| |
| def render(self, test): |
| assert not self._backing_file, 'Placeholder can be used only once' |
| if test.enabled: |
| # cheat and pretend like we're going to pass the data on the |
| # cmdline for test expectation purposes. |
| with contextlib.closing(cStringIO.StringIO()) as output: |
| self.write_encoded_data(output) |
| self._backing_file = output.getvalue() |
| else: # pragma: no cover |
| input_fd, self._backing_file = tempfile.mkstemp(suffix=self.suffix) |
| with os.fdopen(os.dup(input_fd), 'wb') as f: |
| self.write_encoded_data(f) |
| os.close(input_fd) |
| return [self._backing_file] |
| |
| def cleanup(self, test_enabled): |
| assert self._backing_file is not None |
| if not test_enabled: # pragma: no cover |
| try: |
| _rmfile(self._backing_file) |
| except OSError: |
| pass |
| self._backing_file = None |
| |
| def write_encoded_data(self, f): |
| """ Encodes data to be written out, when rendering this placeholder. |
| """ |
| f.write(self.data) |
| |
| |
| class InputTextPlaceholder(InputDataPlaceholder): |
| """ A input placeholder which expects to write out text. |
| """ |
| def __init__(self, data, suffix, name=None): |
| super(InputTextPlaceholder, self).__init__(data, suffix, name=name) |
| assert isinstance(data, basestring) |
| |
| def write_encoded_data(self, f): |
| # Sometimes users give us invalid utf-8 data. They shouldn't, but it does |
| # happen every once and a while. Just ignore it, and replace with �. |
| # We're assuming users only want to write text data out. |
| # self.data can be large, so be careful to do the conversion in chunks |
| # while streaming the data out, instead of requiring a full copy. |
| n = 1 << 16 |
| # This is a generator expression, so this only copies one chunk of |
| # self.data at any one time. |
| chunks = (self.data[i:i + n] for i in xrange(0, len(self.data), n)) |
| decoded = codecs.iterdecode(chunks, 'utf-8', 'replace') |
| for chunk in codecs.iterencode(decoded, 'utf-8'): |
| f.write(chunk) |
| |
| |
| class OutputDataPlaceholder(recipe_util.OutputPlaceholder): |
| def __init__(self, suffix, leak_to, name=None, add_output_log=False): |
| assert add_output_log in (True, False, 'on_failure'), ( |
| 'add_output_log=%r' % add_output_log) |
| self.suffix = suffix |
| self.leak_to = leak_to |
| self.add_output_log = add_output_log |
| self._backing_file = None |
| super(OutputDataPlaceholder, self).__init__(name=name) |
| |
| @property |
| def backing_file(self): |
| return self._backing_file |
| |
| def render(self, test): |
| assert not self._backing_file, 'Placeholder can be used only once' |
| if self.leak_to: |
| self._backing_file = str(self.leak_to) |
| return [self._backing_file] |
| if test.enabled: |
| self._backing_file = '/path/to/tmp/' + self.suffix.lstrip('.') |
| else: # pragma: no cover |
| output_fd, self._backing_file = tempfile.mkstemp(suffix=self.suffix) |
| os.close(output_fd) |
| return [self._backing_file] |
| |
| def result(self, presentation, test): |
| assert self._backing_file |
| ret = None |
| if test.enabled: |
| self._backing_file = None |
| with contextlib.closing(cStringIO.StringIO(test.data or '')) as infile: |
| ret = self.read_decoded_data(infile) |
| else: # pragma: no cover |
| try: |
| with open(self._backing_file, 'rb') as f: |
| ret = self.read_decoded_data(f) |
| finally: |
| if not self.leak_to: |
| _rmfile(self._backing_file) |
| self._backing_file = None |
| |
| if ret is not None and ( |
| self.add_output_log is True or |
| (self.add_output_log == 'on_failure' and |
| presentation.status != 'SUCCESS')): |
| presentation.logs[self.label] = ret.splitlines() |
| |
| return ret |
| |
| def read_decoded_data(self, f): |
| """ Decodes data to be read in, when getting the result of this placeholder. |
| """ |
| return f.read() |
| |
| |
| class OutputTextPlaceholder(OutputDataPlaceholder): |
| """ A output placeholder which expects to write out text. |
| """ |
| def read_decoded_data(self, f): |
| # This ensures that the raw result bytes we got are, in fact, valid utf-8, |
| # replacing invalid bytes with �. Because python2's unicode support is |
| # wonky, we re-encode the now-valid-utf-8 back into a str object so that |
| # users don't need to deal with `unicode` objects. |
| # The file contents can be large, so be careful to do the conversion in |
| # chunks while streaming the data in, instead of requiring a full copy. |
| n = 1 << 16 |
| chunks = iter(lambda: f.read(n), '') |
| decoded = codecs.iterdecode(chunks, 'utf-8', 'replace') |
| return ''.join(codecs.iterencode(decoded, 'utf-8')) |
| |
| class OutputDataDirPlaceholder(recipe_util.OutputPlaceholder): |
| def __init__(self, suffix, leak_to, name=None): |
| self.suffix = suffix |
| self.leak_to = leak_to |
| self._backing_dir = None |
| super(OutputDataDirPlaceholder, self).__init__(name=name) |
| |
| @property |
| def backing_file(self): # pragma: no cover |
| raise ValueError('Output dir placeholders can not be used for stdin, ' |
| 'stdout or stderr') |
| |
| def render(self, test): |
| assert not self._backing_dir, 'Placeholder can be used only once' |
| if self.leak_to: |
| self._backing_dir = str(self.leak_to) |
| if not test.enabled: # pragma: no cover |
| if not os.path.exists(self._backing_dir): |
| os.makedirs(self._backing_dir) |
| else: |
| if not test.enabled: # pragma: no cover |
| self._backing_dir = tempfile.mkdtemp(suffix=self.suffix) |
| else: |
| self._backing_dir = '/path/to/tmp/' + self.suffix |
| |
| return [self._backing_dir] |
| |
| def result(self, presentation, test): |
| assert self._backing_dir |
| if test.enabled: |
| self._backing_dir = None |
| return test.data or {} |
| else: # pragma: no cover |
| try: |
| all_files = {} |
| for dir_path, _, files in os.walk(self._backing_dir): |
| for filename in files: |
| abs_path = os.path.join(dir_path, filename) |
| rel_path = os.path.relpath(abs_path, self._backing_dir) |
| if sys.platform == 'win32': |
| abs_path = ur'\\?\%s' % (abs_path,) |
| with open(abs_path, 'rb') as f: |
| all_files[rel_path] = f.read() |
| return all_files |
| finally: |
| if not self.leak_to: |
| _rmtree(self._backing_dir) |
| self._backing_dir = None |
| |
| |
| class RawIOApi(recipe_api.RecipeApi): |
| @recipe_util.returns_placeholder |
| @staticmethod |
| def input(data, suffix='', name=None): |
| """Returns a Placeholder for use as a step argument. |
| |
| This placeholder can be used to pass data to steps. The recipe engine will |
| dump the 'data' into a file, and pass the filename to the command line |
| argument. |
| |
| data MUST be of type 'str' (not basestring, not unicode). |
| |
| If 'suffix' is not '', it will be used when the engine calls |
| tempfile.mkstemp. |
| |
| See examples/full.py for usage example. |
| """ |
| return InputDataPlaceholder(data, suffix, name=name) |
| |
| @recipe_util.returns_placeholder |
| @staticmethod |
| def input_text(data, suffix='', name=None): |
| """Returns a Placeholder for use as a step argument. |
| |
| data MUST be of type 'str' (not basestring, not unicode). The str is |
| expected to have valid utf-8 data in it. |
| |
| Similar to input(), but ensures that 'data' is valid utf-8 text. Any |
| non-utf-8 characters will be replaced with �. |
| """ |
| return InputTextPlaceholder(data, suffix, name=name) |
| |
| @recipe_util.returns_placeholder |
| @staticmethod |
| def output(suffix='', leak_to=None, name=None, add_output_log=False): |
| """Returns a Placeholder for use as a step argument, or for std{out,err}. |
| |
| If 'leak_to' is None, the placeholder is backed by a temporary file with |
| a suffix 'suffix'. The file is deleted when the step finishes. |
| |
| If 'leak_to' is not None, then it should be a Path and placeholder |
| redirects IO to a file at that path. Once step finishes, the file is |
| NOT deleted (i.e. it's 'leaking'). 'suffix' is ignored in that case. |
| |
| Args: |
| * add_output_log (True|False|'on_failure') - Log a copy of the output |
| to a step link named `name`. If this is 'on_failure', only create this |
| log when the step has a non-SUCCESS status. |
| """ |
| return OutputDataPlaceholder(suffix, leak_to, name=name, |
| add_output_log=add_output_log) |
| |
| @recipe_util.returns_placeholder |
| @staticmethod |
| def output_text(suffix='', leak_to=None, name=None, add_output_log=False): |
| """Returns a Placeholder for use as a step argument, or for std{out,err}. |
| |
| Similar to output(), but uses an OutputTextPlaceholder, which expects utf-8 |
| encoded text. |
| Similar to input(), but tries to decode the resulting data as utf-8 text, |
| replacing any decoding errors with �. |
| |
| Args: |
| * add_output_log (True|False|'on_failure') - Log a copy of the output |
| to a step link named `name`. If this is 'on_failure', only create this |
| log when the step has a non-SUCCESS status. |
| """ |
| return OutputTextPlaceholder(suffix, leak_to, name=name, |
| add_output_log=add_output_log) |
| |
| @recipe_util.returns_placeholder |
| @staticmethod |
| def output_dir(suffix='', leak_to=None, name=None): |
| """Returns a directory Placeholder for use as a step argument. |
| |
| If 'leak_to' is None, the placeholder is backed by a temporary dir with |
| a suffix 'suffix'. The dir is deleted when the step finishes. |
| |
| If 'leak_to' is not None, then it should be a Path and placeholder |
| redirects IO to a dir at that path. Once step finishes, the dir is |
| NOT deleted (i.e. it's 'leaking'). 'suffix' is ignored in that case. |
| """ |
| return OutputDataDirPlaceholder(suffix, leak_to, name=name) |
