Issue #18878: sunau.open now supports the context manager protocol. Based on

patches by Claudiu Popa and R. David Murray.

1 files changed, 23 insertions, 8 deletions

diff --git a/Lib/sunau.py b/Lib/sunau.py
index 010ce23..efdc146 100644
--- a/Lib/sunau.py
+++ b/Lib/sunau.py
@@ -168,6 +168,12 @@ class Au_read:
         if self._file:
             self.close()
 
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *args):
+        self.close()
+
     def initfp(self, file):
         self._file = file
         self._soundpos = 0
@@ -303,6 +309,12 @@ class Au_write:
             self.close()
         self._file = None
 
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *args):
+        self.close()
+
     def initfp(self, file):
         self._file = file
         self._framerate = 0
@@ -410,14 +422,17 @@ class Au_write:
             self._patchheader()
 
     def close(self):
-        self._ensure_header_written()
-        if self._nframeswritten != self._nframes or \
-                  self._datalength != self._datawritten:
-            self._patchheader()
-        self._file.flush()
-        if self._opened and self._file:
-            self._file.close()
-        self._file = None
+        if self._file:
+            try:
+                self._ensure_header_written()
+                if self._nframeswritten != self._nframes or \
+                        self._datalength != self._datawritten:
+                    self._patchheader()
+                self._file.flush()
+            finally:
+                if self._opened and self._file:
+                    self._file.close()
+                self._file = None
 
     #
     # private methods