Merge pull request #139 from us3r64/fix/socket-recv-recv_into

Fix recv, recv_into logic and have same behavior as ESP32SPI

Author: dhalbert

adafruit_wiznet5k/adafruit_wiznet5k_socket.py

@@ -70,16 +70,16 @@ def getdefaulttimeout() -> Optional[float]:
     return _default_socket_timeout
 
 
-def setdefaulttimeout(timeout: Optional[float]) -> None:
+def setdefaulttimeout(_timeout: Optional[float]) -> None:
     """
     Set the default timeout in seconds (float) for new socket objects. When the socket
     module is first imported, the default is None. See settimeout() for possible values
     and their respective meanings.
 
-    :param Optional[float] timeout: The default timeout in seconds or None.
+    :param Optional[float] _timeout: The default timeout in seconds or None.
     """
     global _default_socket_timeout  # pylint: disable=global-statement
-    if timeout is None or (isinstance(timeout, (int, float)) and timeout >= 0):
+    if _timeout is None or (isinstance(timeout, (int, float)) and timeout >= 0):
         _default_socket_timeout = timeout
     else:
         raise ValueError("Timeout must be None, 0.0 or a positive numeric value.")
@@ -450,8 +450,8 @@ def send(self, data: Union[bytes, bytearray]) -> int:
 
         :return int: Number of bytes sent.
         """
-        timeout = 0 if self._timeout is None else self._timeout
-        bytes_sent = _the_interface.socket_write(self._socknum, data, timeout)
+        _timeout = 0 if self._timeout is None else self._timeout
+        bytes_sent = _the_interface.socket_write(self._socknum, data, _timeout)
         gc.collect()
         return bytes_sent
 
@@ -495,21 +495,9 @@ def recv(
 
         :return bytes: Data from the socket.
         """
-        stamp = time.monotonic()
-        while not self._available():
-            if self._timeout and 0 < self._timeout < time.monotonic() - stamp:
-                break
-            time.sleep(0.05)
-        bytes_on_socket = self._available()
-        if not bytes_on_socket:
-            return b""
-        bytes_to_read = min(bytes_on_socket, bufsize)
-        if self._sock_type == SOCK_STREAM:
-            bytes_read = _the_interface.socket_read(self._socknum, bytes_to_read)[1]
-        else:
-            bytes_read = _the_interface.read_udp(self._socknum, bytes_to_read)[1]
-        gc.collect()
-        return bytes(bytes_read)
+        buf = bytearray(bufsize)
+        self.recv_into(buf, bufsize)
+        return bytes(buf)
 
     def _embed_recv(
         self, bufsize: int = 0, flags: int = 0
@@ -568,12 +556,48 @@ def recv_into(self, buffer: bytearray, nbytes: int = 0, flags: int = 0) -> int:
 
         :return int: the number of bytes received
         """
-        if nbytes == 0:
-            nbytes = len(buffer)
-        bytes_received = self.recv(nbytes)
-        nbytes = len(bytes_received)
-        buffer[:nbytes] = bytes_received
-        return nbytes
+        if not 0 <= nbytes <= len(buffer):
+            raise ValueError("nbytes must be 0 to len(buffer)")
+
+        last_read_time = time.monotonic()
+        num_to_read = len(buffer) if nbytes == 0 else nbytes
+        num_read = 0
+        while num_to_read > 0:
+            # we might have read socket data into the self._buffer with:
+            # _readline
+            if len(self._buffer) > 0:
+                bytes_to_read = min(num_to_read, len(self._buffer))
+                buffer[num_read : num_read + bytes_to_read] = self._buffer[
+                    :bytes_to_read
+                ]
+                num_read += bytes_to_read
+                num_to_read -= bytes_to_read
+                self._buffer = self._buffer[bytes_to_read:]
+                # explicitly recheck num_to_read to avoid extra checks
+                continue
+
+            num_avail = self._available()
+            if num_avail > 0:
+                last_read_time = time.monotonic()
+                bytes_to_read = min(num_to_read, num_avail)
+                if self._sock_type == SOCK_STREAM:
+                    bytes_read = _the_interface.socket_read(
+                        self._socknum, bytes_to_read
+                    )[1]
+                else:
+                    bytes_read = _the_interface.read_udp(self._socknum, bytes_to_read)[
+                        1
+                    ]
+                buffer[num_read : num_read + len(bytes_read)] = bytes_read
+                num_read += len(bytes_read)
+                num_to_read -= len(bytes_read)
+            elif num_read > 0:
+                # We got a message, but there are no more bytes to read, so we can stop.
+                break
+            # No bytes yet, or more bytes requested.
+            if self._timeout > 0 and time.monotonic() - last_read_time > self._timeout:
+                raise timeout("timed out")
+        return num_read
 
     @_check_socket_closed
     def recvfrom_into(
@@ -730,3 +754,14 @@ def type(self):
     def proto(self):
         """Socket protocol (always 0x00 in this implementation)."""
         return 0
+
+
+class timeout(TimeoutError):
+    """TimeoutError class. An instance of this error will be raised by recv_into() if
+    the timeout has elapsed and we haven't received any data yet."""
+
+    def __init__(self, msg):
+        super().__init__(msg)
+
+
+# pylint: enable=unused-argument, redefined-builtin, invalid-name