From 1b3452f91dfcd25b21911b86fc9aeab70b5625cd Mon Sep 17 00:00:00 2001 From: Tomas Mlcoch Date: Thu, 15 Jul 2010 14:30:24 +0200 Subject: Improve of mock/disk.py. Add __enter__() and __exit__() methods into TestFile class. Rename "IO" to "TestFile". Add destructor to TestFile and fix open mode "w". Add destructor to TestFile and fix open for write. Catch useless exception. --- tests/mock/disk.py | 17 ++++++++++++++++- 1 file changed, 16 insertions(+), 1 deletion(-) diff --git a/tests/mock/disk.py b/tests/mock/disk.py index 39af5e8c3..1ff2a5c50 100644 --- a/tests/mock/disk.py +++ b/tests/mock/disk.py @@ -36,6 +36,18 @@ class DiskIO(object): def close(self): self.flush() return StringIO.close(self) + + def __del__(self): + try: + self.close() + except (AttributeError): + pass + + def __enter__(self): + return self + + def __exit__(self, *_): + self.close() class Dir(object): pass @@ -65,6 +77,9 @@ class DiskIO(object): content = self.fs.get(path, None) if content == self.Dir: raise IOError("[Errno 21] Is a directory: '%s'" % (path)) + elif mode == "w": + self.fs[path] = "" + f = self.TestFile(self.fs, path, self.fs[path]) elif not content and mode.startswith("w"): self.fs[path] = "" f = self.TestFile(self.fs, path, self.fs[path]) @@ -74,7 +89,7 @@ class DiskIO(object): f = self.TestFile(self.fs, path, content) f.seek(0, os.SEEK_END) else: - f = self.IO(self.fs, path, content) + f = self.TestFile(self.fs, path, content) return f -- cgit