Differentiate recv and recvfrom in sockets

Author: ktnrg45

pyremoteplay/ddp.py

@@ -519,7 +519,7 @@ async def _async_recv_search_msg(
         data = addr = response = None
         if stop_event.is_set():
             return devices
-        response = await sock.recv(0.01)
+        response = await sock.recvfrom(0.01)
         if response is not None:
             data, addr = response
         if data is not None and addr is not None:

pyremoteplay/register.py

@@ -178,7 +178,7 @@ async def _async_regist_init(host: str, host_type: str, timeout: float) -> bool:
         local_addr=(UDP_IP, 0), remote_addr=(host, RP_PORT)
     )
     sock.sendto(data["init"], (host, RP_PORT))
-    response, _ = await sock.recv(timeout=timeout)
+    response = await sock.recv(timeout=timeout)
     success = _check_init(data, response)
     sock.close()
     return success
@@ -212,7 +212,7 @@ async def _async_get_register_info(
     sock = await AsyncTCPSocket.create(remote_addr=(host, RP_PORT))
     sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
     sock.send(b"".join([headers, payload]))
-    response, _ = await sock.recv(timeout)
+    response = await sock.recv(timeout)
     if not response:
         _LOGGER.error("No Register Response Received")
     sock.close()

pyremoteplay/socket.py

@@ -35,10 +35,10 @@ def error_received(self, exc: Exception):
         """Error Received."""
         _LOGGER.error("Socket at: %s received error: %s", self.sock.getsockname(), exc)
 
-    async def recv(
+    async def recvfrom(
         self, timeout: float = None
     ) -> Optional[tuple[bytes, tuple[str, int]]]:
-        """Return received data."""
+        """Return received data and addr."""
         if self.has_callback:
             return None
         try:
@@ -47,17 +47,26 @@ async def recv(
             pass
         return None
 
+    async def recv(self, timeout: float = None) -> Optional[bytes]:
+        """Return received data."""
+        if self.has_callback:
+            return None
+        response = await self.recvfrom(timeout=timeout)
+        if response:
+            return response[0]
+        return None
+
     def sendto(self, data: bytes, *_):
         """Send packet."""
         raise NotImplementedError
 
     def set_callback(
         self, callback: Union[Callable[[bytes, tuple[str, int]], None], None]
     ):
-        """Set callback for datagram received.
+        """Set callback for data received.
 
         Setting this will flush packet received packet queue.
-        :meth:`recv() <pyremoteplay.socket.AsyncUDPProtocol.recv>`
+        :meth:`recv() <pyremoteplay.socket.AsyncBaseSocket.recv>`
         will always return None.
 
         :param callback: callback for data received
@@ -121,7 +130,7 @@ def data_received(self, data: bytes):
             self._packets.put_nowait(item)
 
     def sendto(self, data: bytes, *_):
-        """Send packet."""
+        """Send packet to address."""
         self._transport.write(data)
 
 
@@ -140,21 +149,9 @@ def datagram_received(self, data: bytes, addr: tuple[str, int]):
         else:
             self._packets.put_nowait(item)
 
-    async def recv(
-        self, timeout: float = None
-    ) -> Optional[tuple[bytes, tuple[str, int]]]:
-        """Return received data."""
-        if self.has_callback:
-            return None
-        try:
-            return await asyncio.wait_for(self._packets.get(), timeout=timeout)
-        except (asyncio.TimeoutError, asyncio.CancelledError):
-            pass
-        return None
-
     # pylint: disable=arguments-differ
     def sendto(self, data: bytes, addr: tuple[str, int] = None):
-        """Send packet."""
+        """Send packet to address."""
         self._transport.sendto(data, addr)
 
 
@@ -194,7 +191,7 @@ def sendto(self, data: bytes, addr: tuple[str, int] = None):
         """Send Packet"""
         self._protocol.sendto(data, addr)
 
-    async def recv(self, timeout: float = None):
+    async def recv(self, timeout: float = None) -> Optional[bytes]:
         """Receive a packet."""
 
         packet = await self._protocol.recv(timeout)
@@ -204,6 +201,18 @@ async def recv(self, timeout: float = None):
 
         return packet
 
+    async def recvfrom(
+        self, timeout: float = None
+    ) -> Optional[tuple[bytes, tuple[str, int]]]:
+        """Receive a packet and address."""
+
+        response = await self._protocol.recvfrom(timeout)
+
+        if response is None and self._protocol.closed:
+            raise OSError("Socket is closed")
+
+        return response
+
     def get_extra_info(self, name: str, default: Any = None) -> Any:
         """Return Extra Info."""
         return self._protocol.get_extra_info(name, default)
@@ -215,7 +224,7 @@ def setsockopt(self, __level: int, __optname: int, __value: int | bytes, /):
     def set_callback(
         self, callback: Union[Callable[[bytes, tuple[str, int]], None], None]
     ):
-        """Set callback for datagram received.
+        """Set callback for data received.
 
         Setting this will flush packet received packet queue.
         :meth:`recv() <pyremoteplay.socket.AsyncBaseSocket.recv>`