diff --git a/Lib/socket.py b/Lib/socket.py index 96f8ed0..475a18c 100644 --- a/Lib/socket.py +++ b/Lib/socket.py @@ -48,6 +48,7 @@ import _socket from _socket import * import os, sys, io +from enum import IntEnum try: import errno @@ -60,6 +61,28 @@ EWOULDBLOCK = getattr(errno, 'EWOULDBLOCK', 11) __all__ = ["getfqdn", "create_connection"] __all__.extend(os._get_exports_list(_socket)) +# Set up the socket.AF_* constants as members of an IntEnum for nicer string +# representations. +# Since not all AF_* constants are actually defined by the underlying _socket +# extension on all platforms, this has to be done carefully, by examining the +# constants that did get imported into this module's globals by +# 'from socket import *'. Those constants get replaced by corresponding enum +# members. +# Note that _socket only knows about the integer values. The public interface +# in this module understands the enums and translates them back from integers +# where needed (e.g. .family property of a socket object). +_potential_af_names = ['AF_INET', 'AF_INET6', 'AF_UNIX', 'AF_RDS', 'AF_CAN', + 'AF_LINK', 'AF_UNSPEC', 'AF_AX25', 'AF_IPX', 'AF_NETROM', + 'AF_BRIDGE', 'AD_ATMPVC', 'AF_AAL5', 'AF_X25', 'AF_ROSE', + 'AF_DECnet', 'AF_NETBEUI', 'AF_SECURITY', 'AF_PACKET', + 'AF_NETLINK', 'AF_ROUTE', 'AF_ASH', 'AF_ECONET', + 'AF_ATMSVC', 'AF_SNA', 'AF_IRDA', 'AF_PPPOX', 'AF_KEY', + 'AF_WANPIPE', 'AF_LLC', 'AF_BLUETOOTH', 'AF_SYSTEM'] +_moduledict = globals() +_af_dict = {name: _moduledict[name] for name in _potential_af_names + if name in _moduledict} +AddressFamily = IntEnum('AddressFamily', _af_dict) +_moduledict.update(AddressFamily.__members__) _realsocket = socket @@ -91,6 +114,10 @@ class socket(_socket.socket): __slots__ = ["__weakref__", "_io_refs", "_closed"] def __init__(self, family=AF_INET, type=SOCK_STREAM, proto=0, fileno=None): + # For user code address family values are members of AddressFamily, but + # for the underlying _socket.socket they're just integers. The + # constructor of _socket.socket converts the given family argument to an + # integer automatically. _socket.socket.__init__(self, family, type, proto, fileno) self._io_refs = 0 self._closed = False @@ -229,6 +256,13 @@ class socket(_socket.socket): self._closed = True return super().detach() + @property + def family(self): + """Read-only access to the address family for this socket. + """ + return AddressFamily(super().family) + + def fromfd(fd, family, type, proto=0): """ fromfd(fd, family, type[, proto]) -> socket object diff --git a/Lib/test/test_socket.py b/Lib/test/test_socket.py index 54fb9a1..80e941b 100644 --- a/Lib/test/test_socket.py +++ b/Lib/test/test_socket.py @@ -664,6 +664,11 @@ class GeneralModuleTests(unittest.TestCase): self.assertIn('[closed]', repr(s)) self.assertNotIn('laddr', repr(s)) + def test_str_for_enums(self): + # Make sure that the AF_* constants have enum-like string reprs. + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + self.assertEqual(str(s.family), 'AddressFamily.AF_INET') + def test_weakref(self): s = socket.socket(socket.AF_INET, socket.SOCK_STREAM) p = proxy(s)