diff --git a/mercurial/util.py b/mercurial/util.py --- a/mercurial/util.py +++ b/mercurial/util.py @@ -1516,6 +1516,15 @@ class atomictempfile(object): if safehasattr(self, '_fp'): # constructor actually did something self.discard() + def __enter__(self): + return self + + def __exit__(self, exctype, excvalue, traceback): + if exctype is not None: + self.discard() + else: + self.close() + def makedirs(name, mode=None, notindexed=False): """recursive directory creation with parent mode inheritance diff --git a/tests/test-atomictempfile.py b/tests/test-atomictempfile.py --- a/tests/test-atomictempfile.py +++ b/tests/test-atomictempfile.py @@ -96,6 +96,24 @@ class testatomictempfile(unittest.TestCa self.assertTrue(file.read(), b'foobar\n') file.discard() + def testcontextmanagersuccess(self): + """When the context closes, the file is closed""" + with atomictempfile('foo') as f: + self.assertFalse(os.path.isfile('foo')) + f.write(b'argh\n') + self.assertTrue(os.path.isfile('foo')) + + def testcontextmanagerfailure(self): + """On exception, the file is discarded""" + try: + with atomictempfile('foo') as f: + self.assertFalse(os.path.isfile('foo')) + f.write(b'argh\n') + raise ValueError + except ValueError: + pass + self.assertFalse(os.path.isfile('foo')) + if __name__ == '__main__': import silenttestrunner silenttestrunner.main(__name__)