From ac8b09d853bcf4b27946352024ddaa15a2551152 Mon Sep 17 00:00:00 2001 From: Arash Rouhani Date: Sun, 10 Feb 2019 22:57:28 +0100 Subject: [PATCH] Simplify implementation of temporary_path() The context manager FileSystemTarget.temporary_path() rolled its own manager, it was overcomplicated so i rewrote using contextlib. --- luigi/target.py | 41 ++++++++++++++++------------------------- test/target_test.py | 15 ++++++++------- 2 files changed, 24 insertions(+), 32 deletions(-) diff --git a/luigi/target.py b/luigi/target.py index c9b20fd87a..9ed9665b14 100644 --- a/luigi/target.py +++ b/luigi/target.py @@ -26,6 +26,7 @@ import tempfile import logging import warnings +from contextlib import contextmanager from luigi import six logger = logging.getLogger('luigi-interface') @@ -250,6 +251,7 @@ def remove(self): """ self.fs.remove(self.path) + @contextmanager def temporary_path(self): """ A context manager that enables a reasonably short, general and @@ -275,31 +277,20 @@ def run(self): with self.output().temporary_path() as self.temp_output_path: run_some_external_command(output_path=self.temp_output_path) """ - class _Manager(object): - target = self - - def __init__(self): - num = random.randrange(0, 1e10) - slashless_path = self.target.path.rstrip('/').rstrip("\\") - self._temp_path = '{}-luigi-tmp-{:010}{}'.format( - slashless_path, - num, - self.target._trailing_slash()) - # TODO: os.path doesn't make sense here as it's os-dependent - tmp_dir = os.path.dirname(slashless_path) - if tmp_dir: - self.target.fs.mkdir(tmp_dir, parents=True, raise_if_exists=False) - - def __enter__(self): - return self._temp_path - - def __exit__(self, exc_type, exc_value, traceback): - if exc_type is None: - # There were no exceptions - self.target.fs.rename_dont_move(self._temp_path, self.target.path) - return False # False means we don't suppress the exception - - return _Manager() + num = random.randrange(0, 1e10) + slashless_path = self.path.rstrip('/').rstrip("\\") + _temp_path = '{}-luigi-tmp-{:010}{}'.format( + slashless_path, + num, + self._trailing_slash()) + # TODO: os.path doesn't make sense here as it's os-dependent + tmp_dir = os.path.dirname(slashless_path) + if tmp_dir: + self.fs.mkdir(tmp_dir, parents=True, raise_if_exists=False) + + yield _temp_path + # We won't reach here if there was an user exception. + self.fs.rename_dont_move(_temp_path, self.path) def _touchz(self): with self.open('w'): diff --git a/test/target_test.py b/test/target_test.py index 08486f6d6c..16c4af1601 100644 --- a/test/target_test.py +++ b/test/target_test.py @@ -294,20 +294,21 @@ class MyException(Exception): orig_ex = MyException() try: with target_outer.temporary_path() as tmp_path_outer: - assert 'notreal' in tmp_path_outer + self.assertIn('notreal', tmp_path_outer) with target_inner.temporary_path() as tmp_path_inner: - assert 'blah' in tmp_path_inner + self.assertIn('blah', tmp_path_inner) with target_inner.temporary_path() as tmp_path_inner_2: - assert tmp_path_inner != tmp_path_inner_2 - self.fs.rename_dont_move.assert_called_once_with(tmp_path_inner_2, target_inner.path) + self.assertNotEqual(tmp_path_inner, tmp_path_inner_2) + self.fs.rename_dont_move.assert_called_once_with(tmp_path_inner_2, + target_inner.path) self.fs.rename_dont_move.assert_called_with(tmp_path_inner, target_inner.path) - self.fs.rename_dont_move.call_count == 2 + self.assertEqual(self.fs.rename_dont_move.call_count, 2) raise orig_ex except MyException as ex: - self.fs.rename_dont_move.call_count == 2 - assert ex is orig_ex + self.assertIs(ex, orig_ex) else: assert False + self.assertEqual(self.fs.rename_dont_move.call_count, 2) def test_temporary_path_directory(self): target_slash = self.target_cls('/tmp/dir/')