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

11 files changed, 128 insertions, 106 deletions

diff --git a/Doc/library/http.server.rst b/Doc/library/http.server.rst
index 0bde35b..7ab249a 100644
--- a/Doc/library/http.server.rst
+++ b/Doc/library/http.server.rst
@@ -375,10 +375,9 @@ the current directory::
 
    Handler = http.server.SimpleHTTPRequestHandler
 
-   httpd = socketserver.TCPServer(("", PORT), Handler)
-
-   print("serving at port", PORT)
-   httpd.serve_forever()
+   with socketserver.TCPServer(("", PORT), Handler) as httpd:
+       print("serving at port", PORT)
+       httpd.serve_forever()
 
 .. _http-server-cli:
 
diff --git a/Doc/library/socketserver.rst b/Doc/library/socketserver.rst
index aaaa61e..e148d30 100644
--- a/Doc/library/socketserver.rst
+++ b/Doc/library/socketserver.rst
@@ -52,11 +52,12 @@ handler class by subclassing the :class:`BaseRequestHandler` class and
 overriding its :meth:`~BaseRequestHandler.handle` method;
 this method will process incoming
 requests.  Second, you must instantiate one of the server classes, passing it
-the server's address and the request handler class.  Then call the
+the server's address and the request handler class. It is recommended to use
+the server in a :keyword:`with` statement. Then call the
 :meth:`~BaseServer.handle_request` or
 :meth:`~BaseServer.serve_forever` method of the server object to
 process one or many requests.  Finally, call :meth:`~BaseServer.server_close`
-to close the socket.
+to close the socket (unless you used a :keyword:`with` statement).
 
 When inheriting from :class:`ThreadingMixIn` for threaded connection behavior,
 you should explicitly declare how you want your threads to behave on an abrupt
@@ -353,6 +354,11 @@ Server Objects
       default implementation always returns :const:`True`.
 
 
+   .. versionchanged:: 3.6
+      Support for the :term:`context manager` protocol was added.  Exiting the
+      context manager is equivalent to calling :meth:`server_close`.
+
+
 Request Handler Objects
 -----------------------
 
@@ -433,11 +439,10 @@ This is the server side::
        HOST, PORT = "localhost", 9999
 
        # Create the server, binding to localhost on port 9999
-       server = socketserver.TCPServer((HOST, PORT), MyTCPHandler)
-
-       # Activate the server; this will keep running until you
-       # interrupt the program with Ctrl-C
-       server.serve_forever()
+       with socketserver.TCPServer((HOST, PORT), MyTCPHandler) as server:
+           # Activate the server; this will keep running until you
+           # interrupt the program with Ctrl-C
+           server.serve_forever()
 
 An alternative request handler class that makes use of streams (file-like
 objects that simplify communication by providing the standard file interface)::
@@ -529,8 +534,8 @@ This is the server side::
 
    if __name__ == "__main__":
        HOST, PORT = "localhost", 9999
-       server = socketserver.UDPServer((HOST, PORT), MyUDPHandler)
-       server.serve_forever()
+       with socketserver.UDPServer((HOST, PORT), MyUDPHandler) as server:
+           server.serve_forever()
 
 This is the client side::
 
@@ -592,22 +597,22 @@ An example for the :class:`ThreadingMixIn` class::
        HOST, PORT = "localhost", 0
 
        server = ThreadedTCPServer((HOST, PORT), ThreadedTCPRequestHandler)
-       ip, port = server.server_address
-
-       # Start a thread with the server -- that thread will then start one
-       # more thread for each request
-       server_thread = threading.Thread(target=server.serve_forever)
-       # Exit the server thread when the main thread terminates
-       server_thread.daemon = True
-       server_thread.start()
-       print("Server loop running in thread:", server_thread.name)
-
-       client(ip, port, "Hello World 1")
-       client(ip, port, "Hello World 2")
-       client(ip, port, "Hello World 3")
-
-       server.shutdown()
-       server.server_close()
+       with server:
+           ip, port = server.server_address
+
+           # Start a thread with the server -- that thread will then start one
+           # more thread for each request
+           server_thread = threading.Thread(target=server.serve_forever)
+           # Exit the server thread when the main thread terminates
+           server_thread.daemon = True
+           server_thread.start()
+           print("Server loop running in thread:", server_thread.name)
+
+           client(ip, port, "Hello World 1")
+           client(ip, port, "Hello World 2")
+           client(ip, port, "Hello World 3")
+
+           server.shutdown()
 
 
 The output of the example should look something like this::
diff --git a/Doc/library/wsgiref.rst b/Doc/library/wsgiref.rst
index 71607d6..8d62885 100644
--- a/Doc/library/wsgiref.rst
+++ b/Doc/library/wsgiref.rst
@@ -131,9 +131,9 @@ parameter expect a WSGI-compliant dictionary to be supplied; please see
                  for key, value in environ.items()]
           return ret
 
-      httpd = make_server('', 8000, simple_app)
-      print("Serving on port 8000...")
-      httpd.serve_forever()
+      with make_server('', 8000, simple_app) as httpd:
+          print("Serving on port 8000...")
+          httpd.serve_forever()
 
 
 In addition to the environment functions above, the :mod:`wsgiref.util` module
@@ -283,14 +283,14 @@ request.  (E.g., using the :func:`shift_path_info` function from
 
       from wsgiref.simple_server import make_server, demo_app
 
-      httpd = make_server('', 8000, demo_app)
-      print("Serving HTTP on port 8000...")
+      with make_server('', 8000, demo_app) as httpd:
+          print("Serving HTTP on port 8000...")
 
-      # Respond to requests until process is killed
-      httpd.serve_forever()
+          # Respond to requests until process is killed
+          httpd.serve_forever()
 
-      # Alternative: serve one request, then exit
-      httpd.handle_request()
+          # Alternative: serve one request, then exit
+          httpd.handle_request()
 
 
 .. function:: demo_app(environ, start_response)
@@ -430,9 +430,9 @@ Paste" library.
       # This is the application wrapped in a validator
       validator_app = validator(simple_app)
 
-      httpd = make_server('', 8000, validator_app)
-      print("Listening on port 8000....")
-      httpd.serve_forever()
+      with make_server('', 8000, validator_app) as httpd:
+          print("Listening on port 8000....")
+          httpd.serve_forever()
 
 
 :mod:`wsgiref.handlers` -- server/gateway base classes
@@ -769,8 +769,8 @@ This is a working "Hello World" WSGI application::
        # The returned object is going to be printed
        return [b"Hello World"]
 
-   httpd = make_server('', 8000, hello_world_app)
-   print("Serving on port 8000...")
+   with make_server('', 8000, hello_world_app) as httpd:
+       print("Serving on port 8000...")
 
-   # Serve until process is killed
-   httpd.serve_forever()
+       # Serve until process is killed
+       httpd.serve_forever()
diff --git a/Doc/library/xmlrpc.server.rst b/Doc/library/xmlrpc.server.rst
index 680db41..ca80aab 100644
--- a/Doc/library/xmlrpc.server.rst
+++ b/Doc/library/xmlrpc.server.rst
@@ -147,29 +147,29 @@ Server code::
        rpc_paths = ('/RPC2',)
 
    # Create server
-   server = SimpleXMLRPCServer(("localhost", 8000),
-                               requestHandler=RequestHandler)
-   server.register_introspection_functions()
+   with SimpleXMLRPCServer(("localhost", 8000),
+                           requestHandler=RequestHandler) as server:
+       server.register_introspection_functions()
 
-   # Register pow() function; this will use the value of
-   # pow.__name__ as the name, which is just 'pow'.
-   server.register_function(pow)
+       # Register pow() function; this will use the value of
+       # pow.__name__ as the name, which is just 'pow'.
+       server.register_function(pow)
 
-   # Register a function under a different name
-   def adder_function(x,y):
-       return x + y
-   server.register_function(adder_function, 'add')
+       # Register a function under a different name
+       def adder_function(x,y):
+           return x + y
+       server.register_function(adder_function, 'add')
 
-   # Register an instance; all the methods of the instance are
-   # published as XML-RPC methods (in this case, just 'mul').
-   class MyFuncs:
-       def mul(self, x, y):
-           return x * y
+       # Register an instance; all the methods of the instance are
+       # published as XML-RPC methods (in this case, just 'mul').
+       class MyFuncs:
+           def mul(self, x, y):
+               return x * y
 
-   server.register_instance(MyFuncs())
+       server.register_instance(MyFuncs())
 
-   # Run the server's main loop
-   server.serve_forever()
+       # Run the server's main loop
+       server.serve_forever()
 
 The following client code will call the methods made available by the preceding
 server::
@@ -206,18 +206,17 @@ a server allowing dotted names and registering a multicall function.
             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')
-    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')
+        try:
+            server.serve_forever()
+        except KeyboardInterrupt:
+            print("\nKeyboard interrupt received, exiting.")
+            sys.exit(0)
 
 This ExampleService demo can be invoked from the command line::
 
diff --git a/Doc/whatsnew/3.6.rst b/Doc/whatsnew/3.6.rst
index 8bf2847..ef10ef2 100644
--- a/Doc/whatsnew/3.6.rst
+++ b/Doc/whatsnew/3.6.rst
@@ -259,6 +259,16 @@ you may now specify file paths on top of directories (e.g. zip files).
 (Contributed by Wolfgang Langner in :issue:`26587`).
 
 
+socketserver
+------------
+
+Servers based on the :mod:`socketserver` module, including those
+defined in :mod:`http.server`, :mod:`xmlrpc.server` and
+:mod:`wsgiref.simple_server`, now support the :term:`context manager`
+protocol.
+(Contributed by Aviv Palivoda in :issue:`26404`.)
+
+
 telnetlib
 ---------
 
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)
diff --git a/Misc/NEWS b/Misc/NEWS
index d98c184..a4109ab 100644
--- a/Misc/NEWS
+++ b/Misc/NEWS
@@ -240,6 +240,8 @@ Core and Builtins
 Library
 -------
 
+- Issue #26404: Add context manager to socketserver.  Patch by Aviv Palivoda.
+
 - Issue #26735: Fix :func:`os.urandom` on Solaris 11.3 and newer when reading
   more than 1,024 bytes: call ``getrandom()`` multiple times with a limit of
   1024 bytes per call.