python - 如何在 Python 中使用套接字作为上下文管理器？

Question

似乎做这样的事情是很自然的：

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

但是 Python 没有为套接字实现上下文管理器。我可以轻松地将它用作上下文管理器吗？如果可以，如何使用？

Answer 1

该socket模块相当低级，让您几乎可以直接访问 C 库功能。

你总是可以使用contextlib.contextmanager装饰器来构建你自己的：

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:

或者使用contextlib.closing()来达到同样的效果：

from contextlib import closing

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

但是contextmanager()装饰器让你有机会先用套接字做其他事情。

Python 3.x 确实创建socket()了一个上下文管理器，但文档直到 2016 年 Python 3.5 周期才更新以反映这一点。查看源代码中的socket类__enter__，其中添加和__exit__方法。

Answer 2

套接字模块只是 BSD 套接字接口的一个包装器。它是低级的，并没有真正尝试为您提供方便或易于使用的 Pythonic API。您可能想要使用更高级别的东西。

也就是说，它实际上实现了一个上下文管理器：

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

但是你需要使用 Python 3。

对于 Python 2 兼容性，您可以使用contextlib.

from contextlib import closing
import socket

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

Answer 3

请查看以下片段，用于 TCP 和 UDP 套接字

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()

这样您就可以通过以下方式使用它们：

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)

我还试图强调方法名称中 TCP 和 UDP 之间的“连接”一词的行为和方法的差异。