Socket improvements

- Fixed `socket_available` to just report and not flush available UDP
  data when calling multiple times.
- Fixed socket `__exit__` to only disconnect if a TCP (stream) socket.
- Connect can now accept domain names like CPython versus only IP
  addresses.
- Added socket `send_to`, `recvfrom`, `recv_into`, `recvfrom_into`
  methods for better compatibility with CPython socket.
- Simplified `recv` and `readline` socket methods due to
  `available` doing the right thing for both TCP and UDP.
- Reformatting to make black happy.
- Changes made and tested to make generic socket NTP implementation work:
  https://github.com/askpatrickw/Adafruit_CircuitPython_NTP/blob/raw_ntp/adafruit_ntp.py

Author: xorbit

adafruit_wiznet5k/adafruit_wiznet5k.py

@@ -472,33 +472,20 @@ def socket_available(self, socket_num, sock_type=SNMR_TCP):
             print("* socket_available called with protocol", sock_type)
         assert socket_num <= self.max_sockets, "Provided socket exceeds max_sockets."
 
-        if sock_type == 0x02:
-            # flush by reading remaining data from previous packet
-            while UDP_SOCK["bytes_remaining"] > 0 and self.socket_read(socket_num, 1):
-                if self._debug:
-                    print("Flushing {} bytes".format(UDP_SOCK["bytes_remaining"]))
-                if UDP_SOCK["bytes_remaining"] > 0:
-                    UDP_SOCK["bytes_remaining"] = UDP_SOCK["bytes_remaining"] - 1
-
         res = self._get_rx_rcv_size(socket_num)
 
         if sock_type == SNMR_TCP:
             return res
         if res > 0:
+            if UDP_SOCK["bytes_remaining"]:
+                return UDP_SOCK["bytes_remaining"]
             # parse the udp rx packet
-            ret = 0
             # read the first 8 header bytes
             ret, self._pbuff = self.socket_read(socket_num, 8)
             if ret > 0:
-                UDP_SOCK["remote_ip"] = self._pbuff
-                UDP_SOCK["remote_port"] = self._pbuff[4]
-                UDP_SOCK["remote_port"] = (UDP_SOCK["remote_port"] << 8) + self._pbuff[
-                    5
-                ]
-                UDP_SOCK["bytes_remaining"] = self._pbuff[6]
-                UDP_SOCK["bytes_remaining"] = (
-                    UDP_SOCK["bytes_remaining"] << 8
-                ) + self._pbuff[7]
+                UDP_SOCK["remote_ip"] = self._pbuff[:4]
+                UDP_SOCK["remote_port"] = (self._pbuff[4] << 8) + self._pbuff[5]
+                UDP_SOCK["bytes_remaining"] = (self._pbuff[6] << 8) + self._pbuff[7]
                 ret = UDP_SOCK["bytes_remaining"]
                 return ret
         return 0
@@ -773,17 +760,13 @@ def socket_write(self, socket_num, buffer, timeout=0):
         while (
             self._read_socket(socket_num, REG_SNIR)[0] & SNIR_SEND_OK
         ) != SNIR_SEND_OK:
-            if (
-                self.socket_status(socket_num)[0]
-                in (
-                    SNSR_SOCK_CLOSED,
-                    SNSR_SOCK_TIME_WAIT,
-                    SNSR_SOCK_FIN_WAIT,
-                    SNSR_SOCK_CLOSE_WAIT,
-                    SNSR_SOCK_CLOSING,
-                )
-                or (timeout and time.monotonic() - stamp > timeout)
-            ):
+            if self.socket_status(socket_num)[0] in (
+                SNSR_SOCK_CLOSED,
+                SNSR_SOCK_TIME_WAIT,
+                SNSR_SOCK_FIN_WAIT,
+                SNSR_SOCK_CLOSE_WAIT,
+                SNSR_SOCK_CLOSING,
+            ) or (timeout and time.monotonic() - stamp > timeout):
                 # self.socket_close(socket_num)
                 return 0
             time.sleep(0.01)

adafruit_wiznet5k/adafruit_wiznet5k_socket.py

@@ -84,7 +84,7 @@ def is_ipv4(host):
     return True
 
 
-# pylint: disable=invalid-name
+# pylint: disable=invalid-name, too-many-public-methods
 class socket:
     """A simplified implementation of the Python 'socket' class
     for connecting to a Wiznet5k module.
@@ -112,11 +112,12 @@ def __enter__(self):
         return self
 
     def __exit__(self, exc_type, exc_val, exc_tb):
-        self.disconnect()
-        stamp = time.monotonic()
-        while self.status == adafruit_wiznet5k.SNSR_SOCK_FIN_WAIT:
-            if time.monotonic() - stamp > 1000:
-                raise RuntimeError("Failed to disconnect socket")
+        if self._sock_type == SOCK_STREAM:
+            self.disconnect()
+            stamp = time.monotonic()
+            while self.status == adafruit_wiznet5k.SNSR_SOCK_FIN_WAIT:
+                if time.monotonic() - stamp > 1000:
+                    raise RuntimeError("Failed to disconnect socket")
         self.close()
         stamp = time.monotonic()
         while self.status != adafruit_wiznet5k.SNSR_SOCK_CLOSED:
@@ -216,18 +217,19 @@ def accept(self):
         return client_sock, addr
 
     def connect(self, address, conntype=None):
-        """Connect to a remote socket at address. (The format of address depends
-        on the address family — see above.)
+        """Connect to a remote socket at address.
         :param tuple address: Remote socket as a (host, port) tuple.
-
         """
         assert (
             conntype != 0x03
         ), "Error: SSL/TLS is not currently supported by CircuitPython."
         host, port = address
 
         if hasattr(host, "split"):
-            host = tuple(map(int, host.split(".")))
+            try:
+                host = tuple(map(int, host.split(".")))
+            except ValueError:
+                host = _the_interface.get_host_by_name(host)
         if not _the_interface.socket_connect(
             self.socknum, host, port, conn_mode=self._sock_type
         ):
@@ -238,23 +240,29 @@ def send(self, data):
         """Send data to the socket. The socket must be connected to
         a remote socket.
         :param bytearray data: Desired data to send to the socket.
-
         """
         _the_interface.socket_write(self.socknum, data, self._timeout)
         gc.collect()
 
-    def recv(self, bufsize=0):  # pylint: disable=too-many-branches
+    def sendto(self, data, address):
+        """Send data to the socket. The socket must be connected to
+        a remote socket.
+        :param bytearray data: Desired data to send to the socket.
+        :param tuple address: Remote socket as a (host, port) tuple.
+        """
+        self.connect(address)
+        return self.send(data)
+
+    def recv(self, bufsize=0, flags=0):  # pylint: disable=too-many-branches
         """Reads some bytes from the connected remote address.
         :param int bufsize: Maximum number of bytes to receive.
+        :param int flags: ignored, present for compatibility.
         """
         # print("Socket read", bufsize)
         if bufsize == 0:
             # read everything on the socket
             while True:
-                if self._sock_type == SOCK_STREAM:
-                    avail = self.available()
-                elif self._sock_type == SOCK_DGRAM:
-                    avail = _the_interface.udp_remaining()
+                avail = self.available()
                 if avail:
                     if self._sock_type == SOCK_STREAM:
                         self._buffer += _the_interface.socket_read(self.socknum, avail)[
@@ -275,10 +283,7 @@ def recv(self, bufsize=0):  # pylint: disable=too-many-branches
         received = []
         while to_read > 0:
             # print("Bytes to read:", to_read)
-            if self._sock_type == SOCK_STREAM:
-                avail = self.available()
-            elif self._sock_type == SOCK_DGRAM:
-                avail = _the_interface.udp_remaining()
+            avail = self.available()
             if avail:
                 stamp = time.monotonic()
                 if self._sock_type == SOCK_STREAM:
@@ -305,21 +310,62 @@ def recv(self, bufsize=0):  # pylint: disable=too-many-branches
         gc.collect()
         return ret
 
+    def recvfrom(self, bufsize=0, flags=0):
+        """Reads some bytes from the connected remote address.
+        :param int bufsize: Maximum number of bytes to receive.
+        :param int flags: ignored, present for compatibility.
+        :returns: a tuple (bytes, address) where address is a tuple (ip, port)
+        """
+        return (
+            self.recv(bufsize),
+            (
+                _the_interface.remote_ip(self.socknum),
+                _the_interface.remote_port(self.socknum),
+            ),
+        )
+
+    def recv_into(self, buf, nbytes=0, flags=0):
+        """Reads some bytes from the connected remote address info the provided buffer.
+        :param bytearray buf: Data buffer
+        :param nbytes: Maximum number of bytes to receive
+        :param int flags: ignored, present for compatibility.
+        :returns: the number of bytes received
+        """
+        if nbytes == 0:
+            nbytes = len(buf)
+        ret = self.recv(nbytes)
+        nbytes = len(ret)
+        buf[:nbytes] = ret
+        return nbytes
+
+    def recvfrom_into(self, buf, nbytes=0, flags=0):
+        """Reads some bytes from the connected remote address info the provided buffer.
+        :param bytearray buf: Data buffer
+        :param nbytes: Maximum number of bytes to receive
+        :param int flags: ignored, present for compatibility.
+        :returns a tuple (nbytes, address) where address is a tuple (ip, port)
+        """
+        return (
+            self.recv_into(buf, nbytes),
+            (
+                _the_interface.remote_ip(self.socknum),
+                _the_interface.remote_port(self.socknum),
+            ),
+        )
+
     def readline(self):
         """Attempt to return as many bytes as we can up to \
         but not including '\r\n'.
 
         """
         stamp = time.monotonic()
         while b"\r\n" not in self._buffer:
-            if self._sock_type == SOCK_STREAM:
-                avail = self.available()
-                if avail:
+            avail = self.available()
+            if avail:
+                if self._sock_type == SOCK_STREAM:
                     self._buffer += _the_interface.socket_read(self.socknum, avail)[1]
-            elif self._sock_type == SOCK_DGRAM:
-                avail = _the_interface.udp_remaining()
-                if avail:
-                    self._buffer += _the_interface.read_udp(self.socknum, avail)
+                elif self._sock_type == SOCK_DGRAM:
+                    self._buffer += _the_interface.read_udp(self.socknum, avail)[1]
             if (
                 not avail
                 and self._timeout > 0