Issue #26404: Add context manager to socketserver, by Aviv Palivoda

5 files changed, 38 insertions, 31 deletions

diff --git a/Lib/http/server.py b/Lib/http/server.py
index fbee6a9..c1607b3 100644
--- a/Lib/http/server.py
+++ b/Lib/http/server.py
@@ -1175,16 +1175,14 @@ def test(HandlerClass=BaseHTTPRequestHandler,
     server_address = (bind, port)
 
     HandlerClass.protocol_version = protocol
-    httpd = ServerClass(server_address, HandlerClass)
-
-    sa = httpd.socket.getsockname()
-    print("Serving HTTP on", sa[0], "port", sa[1], "...")
-    try:
-        httpd.serve_forever()
-    except KeyboardInterrupt:
-        print("\nKeyboard interrupt received, exiting.")
-        httpd.server_close()
-        sys.exit(0)
+    with ServerClass(server_address, HandlerClass) as httpd:
+        sa = httpd.socket.getsockname()
+        print("Serving HTTP on", sa[0], "port", sa[1], "...")
+        try:
+            httpd.serve_forever()
+        except KeyboardInterrupt:
+            print("\nKeyboard interrupt received, exiting.")
+            sys.exit(0)
 
 if __name__ == '__main__':
     parser = argparse.ArgumentParser()
diff --git a/Lib/socketserver.py b/Lib/socketserver.py
index 2f39514..3e1f058 100644
--- a/Lib/socketserver.py
+++ b/Lib/socketserver.py
@@ -378,6 +378,12 @@ class BaseServer:
         traceback.print_exc()
         print('-'*40, file=sys.stderr)
 
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *args):
+        self.server_close()
+
 
 class TCPServer(BaseServer):
 
diff --git a/Lib/test/test_socketserver.py b/Lib/test/test_socketserver.py
index 27fe01c..554c106 100644
--- a/Lib/test/test_socketserver.py
+++ b/Lib/test/test_socketserver.py
@@ -104,7 +104,6 @@ class SocketServerTest(unittest.TestCase):
         class MyServer(svrcls):
             def handle_error(self, request, client_address):
                 self.close_request(request)
-                self.server_close()
                 raise
 
         class MyHandler(hdlrbase):
@@ -280,6 +279,12 @@ class SocketServerTest(unittest.TestCase):
                 socketserver.TCPServer((HOST, -1),
                                        socketserver.StreamRequestHandler)
 
+    def test_context_manager(self):
+        with socketserver.TCPServer((HOST, 0),
+                                    socketserver.StreamRequestHandler) as server:
+            pass
+        self.assertEqual(-1, server.socket.fileno())
+
 
 class ErrorHandlerTest(unittest.TestCase):
     """Test that the servers pass normal exceptions from the handler to
diff --git a/Lib/wsgiref/simple_server.py b/Lib/wsgiref/simple_server.py
index 378b316..1807c66 100644
--- a/Lib/wsgiref/simple_server.py
+++ b/Lib/wsgiref/simple_server.py
@@ -156,10 +156,9 @@ def make_server(
 
 
 if __name__ == '__main__':
-    httpd = make_server('', 8000, demo_app)
-    sa = httpd.socket.getsockname()
-    print("Serving HTTP on", sa[0], "port", sa[1], "...")
-    import webbrowser
-    webbrowser.open('http://localhost:8000/xyz?abc')
-    httpd.handle_request()  # serve one request, then exit
-    httpd.server_close()
+    with make_server('', 8000, demo_app) as httpd:
+        sa = httpd.socket.getsockname()
+        print("Serving HTTP on", sa[0], "port", sa[1], "...")
+        import webbrowser
+        webbrowser.open('http://localhost:8000/xyz?abc')
+        httpd.handle_request()  # serve one request, then exit
diff --git a/Lib/xmlrpc/server.py b/Lib/xmlrpc/server.py
index 5b5bf7c..78728f2 100644
--- a/Lib/xmlrpc/server.py
+++ b/Lib/xmlrpc/server.py
@@ -971,16 +971,15 @@ if __name__ == '__main__':
             def getCurrentTime():
                 return datetime.datetime.now()
 
-    server = SimpleXMLRPCServer(("localhost", 8000))
-    server.register_function(pow)
-    server.register_function(lambda x,y: x+y, 'add')
-    server.register_instance(ExampleService(), allow_dotted_names=True)
-    server.register_multicall_functions()
-    print('Serving XML-RPC on localhost port 8000')
-    print('It is advisable to run this example server within a secure, closed network.')
-    try:
-        server.serve_forever()
-    except KeyboardInterrupt:
-        print("\nKeyboard interrupt received, exiting.")
-        server.server_close()
-        sys.exit(0)
+    with SimpleXMLRPCServer(("localhost", 8000)) as server:
+        server.register_function(pow)
+        server.register_function(lambda x,y: x+y, 'add')
+        server.register_instance(ExampleService(), allow_dotted_names=True)
+        server.register_multicall_functions()
+        print('Serving XML-RPC on localhost port 8000')
+        print('It is advisable to run this example server within a secure, closed network.')
+        try:
+            server.serve_forever()
+        except KeyboardInterrupt:
+            print("\nKeyboard interrupt received, exiting.")
+            sys.exit(0)