Search code examples
pythonsocketswith-statementcontextmanager

python `with .. as ..` statement and multiple return values


I am trying to use a python with-statement (a.k.a. a context manager) to ensure that the TCP connection socket created by server_socket.accept() is always closed. However the obvious form does not work because accept() returns multiple values.

Is there a way use a with-statement for functions with multiple return values?

A minimal example is below, I want to use something like the commented code to replace the try/finally block.

#!/usr/bin/env python3

import socket
from socket import socket as Socket

with Socket(socket.AF_INET, socket.SOCK_STREAM) as server_socket:

        server_socket.bind(('', 8011))
        server_socket.listen(1)
        print("server ready")

        while True:

            # with server_socket.accept() as (connection_socket, _):
            #     request = connection_socket.recv(1024).decode('ascii')
            #     reply = request.upper()
            #     connection_socket.send(reply.encode('ascii'))

            try:
                connection_socket, _ = server_socket.accept()
                request = connection_socket.recv(1024).decode('ascii')
                reply = request.upper()
                connection_socket.send(reply.encode('ascii'))

            finally:
                connection_socket.close()

The error message when using the commented with-statement is:

Traceback (most recent call last):
  File "./test.py", line 26, in <module>
    with server_socket.accept() as (connection_socket, _):
AttributeError: __exit__

Presumably this is because a tuple does not have the __exit__ attribute required for with.


Solution

  • The return value of socket.socket has a built-in context manager that implements __exit__ and __enter__. The tuple returned by accept does not, but you can break out the client socket and address from the tuple and use with on the client socket alone:

    import socket
    
    with socket.socket() as s:
        s.bind(('',8000))
        s.listen(1)
        c,a = s.accept(s)
        with c:
            while True:
                data = c.recv(1024)
                if not data: break
                c.sendall(data)