Sockets facelift. APIs that could return binary data (e.g. aton() and

recv()) now return bytes, not str or str8.  The socket.py code is
redone; it now subclasses _socket.socket and instead of having its own
_fileobject for makefile(), it uses io.SocketIO.  Some stuff in io.py
was moved around to make this work.  (I really need to rethink my
policy regarding readline() and read(-1) on raw files; and readline()
on buffered files ought to use peeking().  Later.)

4 files changed, 148 insertions, 360 deletions

diff --git a/Lib/io.py b/Lib/io.py
index 5f503c2..9cbc11c 100644
--- a/Lib/io.py
+++ b/Lib/io.py
@@ -295,6 +295,22 @@ class IOBase:
         """
         return False
 
+    ### Readline ###
+
+    def readline(self, sizehint: int = -1) -> bytes:
+        """For backwards compatibility, a (slow) readline()."""
+        if sizehint is None:
+            sizehint = -1
+        res = b""
+        while sizehint < 0 or len(res) < sizehint:
+            b = self.read(1)
+            if not b:
+                break
+            res += b
+            if b == b"\n":
+                break
+        return res
+
 
 class RawIOBase(IOBase):
 
@@ -366,7 +382,6 @@ class SocketIO(RawIOBase):
     """Raw I/O implementation for stream sockets."""
 
     # XXX More docs
-    # XXX Hook this up to socket.py
 
     def __init__(self, sock, mode):
         assert mode in ("r", "w", "rw")
@@ -377,13 +392,32 @@ class SocketIO(RawIOBase):
     def readinto(self, b):
         return self._sock.recv_into(b)
 
+    def read(self, n: int = None) -> bytes:
+        """read(n: int) -> bytes.  Read and return up to n bytes.
+
+        Returns an empty bytes array on EOF, or None if the object is
+        set not to block and has no data to read.
+        """
+        if n is None:
+            n = -1
+        if n >= 0:
+            return RawIOBase.read(self, n)
+        # Support reading until the end.
+        # XXX Why doesn't RawIOBase support this?
+        data = b""
+        while True:
+            more = RawIOBase.read(self, DEFAULT_BUFFER_SIZE)
+            if not more:
+                break
+            data += more
+        return data
+
     def write(self, b):
         return self._sock.send(b)
 
     def close(self):
         if not self.closed:
-            RawIOBase.close()
-            self._sock.close()
+            RawIOBase.close(self)
 
     def readable(self):
         return "r" in self._mode
@@ -450,20 +484,6 @@ class BufferedIOBase(IOBase):
         b[:n] = data
         return n
 
-    def readline(self, sizehint: int = -1) -> bytes:
-        """For backwards compatibility, a (slow) readline()."""
-        if sizehint is None:
-            sizehint = -1
-        res = b""
-        while sizehint < 0 or len(res) < sizehint:
-            b = self.read(1)
-            if not b:
-                break
-            res += b
-            if b == b"\n":
-                break
-        return res
-
     def write(self, b: bytes) -> int:
         """write(b: bytes) -> int.  Write the given buffer to the IO stream.
 
diff --git a/Lib/socket.py b/Lib/socket.py
index 8dd2383..03cdc65 100644
--- a/Lib/socket.py
+++ b/Lib/socket.py
@@ -54,7 +54,7 @@ try:
 except ImportError:
     pass
 
-import os, sys
+import os, sys, io
 
 try:
     from errno import EBADF
@@ -66,14 +66,6 @@ __all__.extend(os._get_exports_list(_socket))
 if _have_ssl:
     __all__.extend(os._get_exports_list(_ssl))
 
-_realsocket = socket
-if _have_ssl:
-    _realssl = ssl
-    def ssl(sock, keyfile=None, certfile=None):
-        if hasattr(sock, "_sock"):
-            sock = sock._sock
-        return _realssl(sock, keyfile, certfile)
-
 # WSA error codes
 if sys.platform.lower().startswith("win"):
     errorTab = {}
@@ -95,6 +87,99 @@ if sys.platform.lower().startswith("win"):
     __all__.append("errorTab")
 
 
+_os_has_dup = hasattr(os, "dup")
+if _os_has_dup:
+    def fromfd(fd, family=AF_INET, type=SOCK_STREAM, proto=0):
+        nfd = os.dup(fd)
+        return socket(family, type, proto, fileno=nfd)
+
+
+class socket(_socket.socket):
+
+    """A subclass of _socket.socket adding the makefile() method."""
+
+    __slots__ = ["__weakref__"]
+    if not _os_has_dup:
+        __slots__.append("_base")
+
+    def __repr__(self):
+        """Wrap __repr__() to reveal the real class name."""
+        s = _socket.socket.__repr__(self)
+        if s.startswith("<socket object"):
+            s = "<%s.%s%s" % (self.__class__.__module__,
+                              self.__class__.__name__,
+                              s[7:])
+        return s
+
+    def accept(self):
+        """Wrap accept() to give the connection the right type."""
+        conn, addr = _socket.socket.accept(self)
+        fd = conn.fileno()
+        nfd = fd
+        if _os_has_dup:
+            nfd = os.dup(fd)
+        wrapper = socket(self.family, self.type, self.proto, fileno=nfd)
+        if fd == nfd:
+            wrapper._base = conn  # Keep the base alive
+        else:
+            conn.close()
+        return wrapper, addr
+
+    if not _os_has_dup:
+        def close(self):
+            """Wrap close() to close the _base as well."""
+            _socket.socket.close(self)
+            base = getattr(self, "_base", None)
+            if base is not None:
+                base.close()
+
+    def makefile(self, mode="r", buffering=None, *,
+                 encoding=None, newline=None):
+        """Return an I/O stream connected to the socket.
+
+        The arguments are as for io.open() after the filename,
+        except the only mode characters supported are 'r', 'w' and 'b'.
+        The semantics are similar too.  (XXX refactor to share code?)
+        """
+        for c in mode:
+            if c not in {"r", "w", "b"}:
+                raise ValueError("invalid mode %r (only r, w, b allowed)")
+        writing = "w" in mode
+        reading = "r" in mode or not writing
+        assert reading or writing
+        binary = "b" in mode
+        rawmode = ""
+        if reading:
+            rawmode += "r"
+        if writing:
+            rawmode += "w"
+        raw = io.SocketIO(self, rawmode)
+        if buffering is None:
+            buffering = -1
+        if buffering < 0:
+            buffering = io.DEFAULT_BUFFER_SIZE
+        if buffering == 0:
+            if not binary:
+                raise ValueError("unbuffered streams must be binary")
+            raw.name = self.fileno()
+            raw.mode = mode
+            return raw
+        if reading and writing:
+            buffer = io.BufferedRWPair(raw, raw, buffering)
+        elif reading:
+            buffer = io.BufferedReader(raw, buffering)
+        else:
+            assert writing
+            buffer = io.BufferedWriter(raw, buffering)
+        if binary:
+            buffer.name = self.fileno()
+            buffer.mode = mode
+            return buffer
+        text = io.TextIOWrapper(buffer, encoding, newline)
+        text.name = self.fileno()
+        self.mode = mode
+        return text
+
 
 def getfqdn(name=''):
     """Get fully qualified domain name from name.
@@ -122,298 +207,6 @@ def getfqdn(name=''):
     return name
 
 
-_socketmethods = (
-    'bind', 'connect', 'connect_ex', 'fileno', 'listen',
-    'getpeername', 'getsockname', 'getsockopt', 'setsockopt',
-    'sendall', 'setblocking',
-    'settimeout', 'gettimeout', 'shutdown')
-
-if sys.platform == "riscos":
-    _socketmethods = _socketmethods + ('sleeptaskw',)
-
-# All the method names that must be delegated to either the real socket
-# object or the _closedsocket object.
-_delegate_methods = ("recv", "recvfrom", "recv_into", "recvfrom_into",
-                     "send", "sendto")
-
-class _closedsocket(object):
-    __slots__ = []
-    def _dummy(*args):
-        raise error(EBADF, 'Bad file descriptor')
-    # All _delegate_methods must also be initialized here.
-    send = recv = recv_into = sendto = recvfrom = recvfrom_into = _dummy
-    __getattr__ = _dummy
-
-class _socketobject(object):
-
-    __doc__ = _realsocket.__doc__
-
-    __slots__ = ["_sock", "__weakref__"] + list(_delegate_methods)
-
-    def __init__(self, family=AF_INET, type=SOCK_STREAM, proto=0, _sock=None):
-        if _sock is None:
-            _sock = _realsocket(family, type, proto)
-        self._sock = _sock
-        for method in _delegate_methods:
-            setattr(self, method, getattr(_sock, method))
-
-    def close(self):
-        self._sock = _closedsocket()
-        dummy = self._sock._dummy
-        for method in _delegate_methods:
-            setattr(self, method, dummy)
-    close.__doc__ = _realsocket.close.__doc__
-
-    def accept(self):
-        sock, addr = self._sock.accept()
-        return _socketobject(_sock=sock), addr
-    accept.__doc__ = _realsocket.accept.__doc__
-
-    def dup(self):
-        """dup() -> socket object
-
-        Return a new socket object connected to the same system resource."""
-        return _socketobject(_sock=self._sock)
-
-    def makefile(self, mode='r', bufsize=-1):
-        """makefile([mode[, bufsize]]) -> file object
-
-        Return a regular file object corresponding to the socket.  The mode
-        and bufsize arguments are as for the built-in open() function."""
-        return _fileobject(self._sock, mode, bufsize)
-
-    family = property(lambda self: self._sock.family, doc="the socket family")
-    type = property(lambda self: self._sock.type, doc="the socket type")
-    proto = property(lambda self: self._sock.proto, doc="the socket protocol")
-
-    _s = ("def %s(self, *args): return self._sock.%s(*args)\n\n"
-          "%s.__doc__ = _realsocket.%s.__doc__\n")
-    for _m in _socketmethods:
-        exec(_s % (_m, _m, _m, _m))
-    del _m, _s
-
-socket = SocketType = _socketobject
-
-class _fileobject(object):
-    """Faux file object attached to a socket object."""
-
-    default_bufsize = 8192
-    name = "<socket>"
-
-    __slots__ = ["mode", "bufsize",
-                 # "closed" is a property, see below
-                 "_sock", "_rbufsize", "_wbufsize", "_rbuf", "_wbuf",
-                 "_close"]
-
-    def __init__(self, sock, mode='rb', bufsize=-1, close=False):
-        self._sock = sock
-        self.mode = mode # Not actually used in this version
-        if bufsize < 0:
-            bufsize = self.default_bufsize
-        self.bufsize = bufsize
-        if bufsize == 0:
-            self._rbufsize = 1
-        elif bufsize == 1:
-            self._rbufsize = self.default_bufsize
-        else:
-            self._rbufsize = bufsize
-        self._wbufsize = bufsize
-        self._rbuf = "" # A string
-        self._wbuf = [] # A list of strings
-        self._close = close
-
-    def _getclosed(self):
-        return self._sock is None
-    closed = property(_getclosed, doc="True if the file is closed")
-
-    def close(self):
-        try:
-            if self._sock:
-                self.flush()
-        finally:
-            if self._close:
-                self._sock.close()
-            self._sock = None
-
-    def __del__(self):
-        try:
-            self.close()
-        except:
-            # close() may fail if __init__ didn't complete
-            pass
-
-    def flush(self):
-        if self._wbuf:
-            buffer = "".join(self._wbuf)
-            self._wbuf = []
-            self._sock.sendall(buffer)
-
-    def fileno(self):
-        return self._sock.fileno()
-
-    def write(self, data):
-        data = str(data) # XXX Should really reject non-string non-buffers
-        if not data:
-            return
-        self._wbuf.append(data)
-        if (self._wbufsize == 0 or
-            self._wbufsize == 1 and '\n' in data or
-            self._get_wbuf_len() >= self._wbufsize):
-            self.flush()
-
-    def writelines(self, list):
-        # XXX We could do better here for very long lists
-        # XXX Should really reject non-string non-buffers
-        self._wbuf.extend(filter(None, map(str, list)))
-        if (self._wbufsize <= 1 or
-            self._get_wbuf_len() >= self._wbufsize):
-            self.flush()
-
-    def _get_wbuf_len(self):
-        buf_len = 0
-        for x in self._wbuf:
-            buf_len += len(x)
-        return buf_len
-
-    def read(self, size=-1):
-        data = self._rbuf
-        if size < 0:
-            # Read until EOF
-            buffers = []
-            if data:
-                buffers.append(data)
-            self._rbuf = ""
-            if self._rbufsize <= 1:
-                recv_size = self.default_bufsize
-            else:
-                recv_size = self._rbufsize
-            while True:
-                data = self._sock.recv(recv_size)
-                if not data:
-                    break
-                buffers.append(data)
-            return "".join(buffers)
-        else:
-            # Read until size bytes or EOF seen, whichever comes first
-            buf_len = len(data)
-            if buf_len >= size:
-                self._rbuf = data[size:]
-                return data[:size]
-            buffers = []
-            if data:
-                buffers.append(data)
-            self._rbuf = ""
-            while True:
-                left = size - buf_len
-                recv_size = max(self._rbufsize, left)
-                data = self._sock.recv(recv_size)
-                if not data:
-                    break
-                buffers.append(data)
-                n = len(data)
-                if n >= left:
-                    self._rbuf = data[left:]
-                    buffers[-1] = data[:left]
-                    break
-                buf_len += n
-            return "".join(buffers)
-
-    def readline(self, size=-1):
-        data = self._rbuf
-        if size < 0:
-            # Read until \n or EOF, whichever comes first
-            if self._rbufsize <= 1:
-                # Speed up unbuffered case
-                assert data == ""
-                buffers = []
-                recv = self._sock.recv
-                while data != "\n":
-                    data = recv(1)
-                    if not data:
-                        break
-                    buffers.append(data)
-                return "".join(buffers)
-            nl = data.find('\n')
-            if nl >= 0:
-                nl += 1
-                self._rbuf = data[nl:]
-                return data[:nl]
-            buffers = []
-            if data:
-                buffers.append(data)
-            self._rbuf = ""
-            while True:
-                data = self._sock.recv(self._rbufsize)
-                if not data:
-                    break
-                buffers.append(data)
-                nl = data.find('\n')
-                if nl >= 0:
-                    nl += 1
-                    self._rbuf = data[nl:]
-                    buffers[-1] = data[:nl]
-                    break
-            return "".join(buffers)
-        else:
-            # Read until size bytes or \n or EOF seen, whichever comes first
-            nl = data.find('\n', 0, size)
-            if nl >= 0:
-                nl += 1
-                self._rbuf = data[nl:]
-                return data[:nl]
-            buf_len = len(data)
-            if buf_len >= size:
-                self._rbuf = data[size:]
-                return data[:size]
-            buffers = []
-            if data:
-                buffers.append(data)
-            self._rbuf = ""
-            while True:
-                data = self._sock.recv(self._rbufsize)
-                if not data:
-                    break
-                buffers.append(data)
-                left = size - buf_len
-                nl = data.find('\n', 0, left)
-                if nl >= 0:
-                    nl += 1
-                    self._rbuf = data[nl:]
-                    buffers[-1] = data[:nl]
-                    break
-                n = len(data)
-                if n >= left:
-                    self._rbuf = data[left:]
-                    buffers[-1] = data[:left]
-                    break
-                buf_len += n
-            return "".join(buffers)
-
-    def readlines(self, sizehint=0):
-        total = 0
-        list = []
-        while True:
-            line = self.readline()
-            if not line:
-                break
-            list.append(line)
-            total += len(line)
-            if sizehint and total >= sizehint:
-                break
-        return list
-
-    # Iterator protocols
-
-    def __iter__(self):
-        return self
-
-    def __next__(self):
-        line = self.readline()
-        if not line:
-            raise StopIteration
-        return line
-
-
 def create_connection(address, timeout=None):
     """Connect to address (host, port) with an optional timeout.
 
diff --git a/Lib/test/test_socket.py b/Lib/test/test_socket.py
index 350aded..5158378 100644
--- a/Lib/test/test_socket.py
+++ b/Lib/test/test_socket.py
@@ -14,7 +14,7 @@ import signal
 
 PORT = 50007
 HOST = 'localhost'
-MSG = 'Michael Gilfix was here\n'
+MSG = b'Michael Gilfix was here\n'
 
 class SocketTCPTest(unittest.TestCase):
 
@@ -542,16 +542,16 @@ class BasicTCPTest(SocketConnectedTest):
 
     def testSendAll(self):
         # Testing sendall() with a 2048 byte string over TCP
-        msg = ''
+        msg = b''
         while 1:
             read = self.cli_conn.recv(1024)
             if not read:
                 break
             msg += read
-        self.assertEqual(msg, 'f' * 2048)
+        self.assertEqual(msg, b'f' * 2048)
 
     def _testSendAll(self):
-        big_chunk = 'f' * 2048
+        big_chunk = b'f' * 2048
         self.serv_conn.sendall(big_chunk)
 
     def testFromFd(self):
@@ -612,7 +612,7 @@ class TCPCloserTest(ThreadedTCPSocketTest):
         sd = self.cli
         read, write, err = select.select([sd], [], [], 1.0)
         self.assertEqual(read, [sd])
-        self.assertEqual(sd.recv(1), '')
+        self.assertEqual(sd.recv(1), b'')
 
     def _testClose(self):
         self.cli.connect((HOST, PORT))
@@ -754,7 +754,7 @@ class FileObjectClassTestCase(SocketConnectedTest):
 
     def testUnbufferedRead(self):
         # Performing unbuffered file read test
-        buf = ''
+        buf = b''
         while 1:
             char = self.serv_file.read(1)
             if not char:
@@ -796,14 +796,14 @@ class UnbufferedFileObjectClassTestCase(FileObjectClassTestCase):
     def testUnbufferedReadline(self):
         # Read a line, create a new file object, read another line with it
         line = self.serv_file.readline() # first line
-        self.assertEqual(line, "A. " + MSG) # first line
+        self.assertEqual(line, b"A. " + MSG) # first line
         self.serv_file = self.cli_conn.makefile('rb', 0)
         line = self.serv_file.readline() # second line
-        self.assertEqual(line, "B. " + MSG) # second line
+        self.assertEqual(line, b"B. " + MSG) # second line
 
     def _testUnbufferedReadline(self):
-        self.cli_file.write("A. " + MSG)
-        self.cli_file.write("B. " + MSG)
+        self.cli_file.write(b"A. " + MSG)
+        self.cli_file.write(b"B. " + MSG)
         self.cli_file.flush()
 
 class LineBufferedFileObjectClassTestCase(FileObjectClassTestCase):
@@ -818,6 +818,7 @@ class SmallBufferedFileObjectClassTestCase(FileObjectClassTestCase):
 
 class NetworkConnectionTest(object):
     """Prove network connection."""
+
     def clientSetUp(self):
         self.cli = socket.create_connection((HOST, PORT))
         self.serv_conn = self.cli
@@ -827,6 +828,7 @@ class BasicTCPTest2(NetworkConnectionTest, BasicTCPTest):
     """
 
 class NetworkConnectionNoServer(unittest.TestCase):
+
     def testWithoutServer(self):
         self.failUnlessRaises(socket.error, lambda: socket.create_connection((HOST, PORT)))
 
@@ -895,43 +897,19 @@ class NetworkConnectionBehaviourTest(SocketTCPTest, ThreadableTest):
     def testInsideTimeout(self):
         conn, addr = self.serv.accept()
         time.sleep(3)
-        conn.send("done!")
+        conn.send(b"done!")
     testOutsideTimeout = testInsideTimeout
 
     def _testInsideTimeout(self):
         self.cli = sock = socket.create_connection((HOST, PORT))
         data = sock.recv(5)
-        self.assertEqual(data, "done!")
+        self.assertEqual(data, b"done!")
 
     def _testOutsideTimeout(self):
         self.cli = sock = socket.create_connection((HOST, PORT), timeout=1)
         self.failUnlessRaises(socket.timeout, lambda: sock.recv(5))
 
 
-class Urllib2FileobjectTest(unittest.TestCase):
-
-    # urllib2.HTTPHandler has "borrowed" socket._fileobject, and requires that
-    # it close the socket if the close c'tor argument is true
-
-    def testClose(self):
-        class MockSocket:
-            closed = False
-            def flush(self): pass
-            def close(self): self.closed = True
-
-        # must not close unless we request it: the original use of _fileobject
-        # by module socket requires that the underlying socket not be closed until
-        # the _socketobject that created the _fileobject is closed
-        s = MockSocket()
-        f = socket._fileobject(s)
-        f.close()
-        self.assert_(not s.closed)
-
-        s = MockSocket()
-        f = socket._fileobject(s, close=True)
-        f.close()
-        self.assert_(s.closed)
-
 class TCPTimeoutTest(SocketTCPTest):
 
     def testTCPTimeout(self):
@@ -1055,7 +1033,7 @@ class BufferIOTest(SocketConnectedTest):
         buf = b" "*1024
         nbytes = self.cli_conn.recv_into(buf)
         self.assertEqual(nbytes, len(MSG))
-        msg = str(buf[:len(MSG)])
+        msg = buf[:len(MSG)]
         self.assertEqual(msg, MSG)
 
     def _testRecvInto(self):
@@ -1066,7 +1044,7 @@ class BufferIOTest(SocketConnectedTest):
         buf = b" "*1024
         nbytes, addr = self.cli_conn.recvfrom_into(buf)
         self.assertEqual(nbytes, len(MSG))
-        msg = str(buf[:len(MSG)])
+        msg = buf[:len(MSG)]
         self.assertEqual(msg, MSG)
 
     def _testRecvFromInto(self):
@@ -1085,7 +1063,6 @@ def test_main():
         UnbufferedFileObjectClassTestCase,
         LineBufferedFileObjectClassTestCase,
         SmallBufferedFileObjectClassTestCase,
-        Urllib2FileobjectTest,
         NetworkConnectionNoServer,
         NetworkConnectionAttributesTest,
         NetworkConnectionBehaviourTest,
diff --git a/Lib/urllib2.py b/Lib/urllib2.py
index 284c921..4127485 100644
--- a/Lib/urllib2.py
+++ b/Lib/urllib2.py
@@ -1085,10 +1085,8 @@ class AbstractHTTPHandler(BaseHandler):
         # to read().  This weird wrapping allows the returned object to
         # have readline() and readlines() methods.
 
-        # XXX It might be better to extract the read buffering code
-        # out of socket._fileobject() and into a base class.
-
         r.recv = r.read
+        # XXX socket._fileobject is gone; use some class from io.py instead
         fp = socket._fileobject(r, close=True)
 
         resp = addinfourl(fp, r.msg, req.get_full_url())