How to use socket in Python as a context manager?

Question

It seems like it would be only natural to do something like:

with socket(socket.AF_INET, socket.SOCK_DGRAM) as s:

but Python doesn't implement a context manager for socket. Can I easily use it as a context manager, and if so, how?

Answer 1

Please have a look on following snippets, for both TCP and UDP sockets

import socket
from contextlib import contextmanager


@contextmanager
def tcp_connection_to(*args, **kwargs):
    s = socket.create_connection(*args, **kwargs)
    yield s
    s.close()


@contextmanager
def udp_connection():
    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    yield s
    s.close()

So that you can use them in following way:

MY_SERVER = ('localhost', 5000)   # Yes, we need tuple here
some_data = bytes("Hello.")

with tcp_connection_to(MY_SERVER) as conn:
    conn.send(some_data)

with udp_connection() as conn:
    conn.sendto(some_data, MY_SERVER)

I've also tried to emphasise the difference in behaviour and approach to term 'connection' between TCP and UDP in method names.

Answer 2

The socket module is just a wrapper around the BSD socket interface. It's low-level, and does not really attempt to provide you with a handy or easy to use Pythonic API. You may want to use something higher-level.

That said, it does in fact implement a context manager:

>>> with socket.socket() as s:
...   print(s)
... 
<socket.socket object, fd=3, family=2, type=1, proto=0>

But you need to use Python 3.

For Python 2 compatibility you can use contextlib.

from contextlib import closing
import socket

with closing(socket.socket()) as s:
    print s

Answer 3

The socket module is fairly low-level, giving you almost direct access to the C library functionality.

You can always use the contextlib.contextmanager decorator to build your own:

import socket
from contextlib import contextmanager

@contextmanager
def socketcontext(*args, **kw):
    s = socket.socket(*args, **kw)
    try:
        yield s
    finally:
        s.close()

with socketcontext(socket.AF_INET, socket.SOCK_DGRAM) as s:

or use contextlib.closing() to achieve the same effect:

from contextlib import closing

with closing(socket.socket(socket.AF_INET, socket.SOCK_DGRAM)) as s:

but the contextmanager() decorator gives you the opportunity to do other things with the socket first.

Python 3.x does make socket() a context manager, even though the documentation fails to mention that. See the socket class in the source code, which adds __enter__ and __exit__ methods.