Issue 9795: adds context manager protocol to nntplib.NNTP class so that it can used with the 'with' statement.

2 files changed, 35 insertions, 0 deletions

diff --git a/Lib/nntplib.py b/Lib/nntplib.py
index 70a75b6..ba70192 100644
--- a/Lib/nntplib.py
+++ b/Lib/nntplib.py
@@ -346,6 +346,20 @@ class _NNTPBase:
         # Log in and encryption setup order is left to subclasses.
         self.authenticated = False
 
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *args):
+        is_connected = lambda: hasattr(self, "file")
+        if is_connected():
+            try:
+                self.quit()
+            except (socket.error, EOFError):
+                pass
+            finally:
+                if is_connected():
+                    self._close()
+
     def getwelcome(self):
         """Get the welcome message from the server
         (this is read and squirreled away by __init__()).
diff --git a/Lib/test/test_nntplib.py b/Lib/test/test_nntplib.py
index e463e52..ec790ad 100644
--- a/Lib/test/test_nntplib.py
+++ b/Lib/test/test_nntplib.py
@@ -1,4 +1,5 @@
 import io
+import socket
 import datetime
 import textwrap
 import unittest
@@ -252,6 +253,26 @@ class NetworkedNNTPTestsMixin:
             # value
             setattr(cls, name, wrap_meth(meth))
 
+    def test_with_statement(self):
+        def is_connected():
+            if not hasattr(server, 'file'):
+                return False
+            try:
+                server.help()
+            except (socket.error, EOFError):
+                return False
+            return True
+
+        with self.NNTP_CLASS(self.NNTP_HOST, timeout=TIMEOUT, usenetrc=False) as server:
+            self.assertTrue(is_connected())
+            self.assertTrue(server.help())
+        self.assertFalse(is_connected())
+
+        with self.NNTP_CLASS(self.NNTP_HOST, timeout=TIMEOUT, usenetrc=False) as server:
+            server.quit()
+        self.assertFalse(is_connected())
+
+
 NetworkedNNTPTestsMixin.wrap_methods()