Merge pull request #3 from Amund211/main

.gitignore

@@ -1,5 +1,6 @@
 __pycache__/
 dist/
+.mypy_cache/
 *.egg-info/
 *.pyc
 

discordrp/__init__.py

@@ -2,10 +2,11 @@
 A lightweight and safe module for creating custom rich presences on Discord.
 """
 
-from .presence import Presence
+from .presence import Presence, PresenceError
 
 __title__ = 'discord-rich-presence'
 __author__ = 'TenType'
 __copyright__ = 'Copyright 2022 TenType'
 __license__ = 'MIT'
 __version__ = '1.0.2'
+__all__ = ("Presence", "PresenceError")

discordrp/presence.py

@@ -4,14 +4,17 @@
 import struct
 import sys
 
+from abc import ABC, abstractmethod
 from enum import IntEnum
-from typing import Any
+from typing import Any, Optional, cast
+from types import TracebackType
 from uuid import uuid4
 
 class OpCode(IntEnum):
     """
     A list of valid opcodes that can be sent in packets to Discord.
     """
+
     HANDSHAKE = 0
     FRAME = 1
     CLOSE = 2
@@ -34,16 +37,14 @@ class Presence:
 
     def __init__(self, client_id: str):
         self.client_id = client_id
-        self._platform = sys.platform
-        self._socket = None
 
         # Connect to Discord IPC
-        self._connect()
-
+        self._socket: _Socket = _WindowsSocket() if sys.platform == 'win32' else _UnixSocket()
+    
         # Send a handshake request
         self._handshake()
 
-    def set(self, activity):
+    def set(self, activity: Optional[dict[str, Any]]) -> None:
         """
         Sends an activity payload to Discord.
         :param activity: A dictionary of this format:
@@ -86,100 +87,128 @@ def set(self, activity):
         }
         self._send(payload, OpCode.FRAME)
 
-    def clear(self):
+    def clear(self) -> None:
         """
         Clears the current activity.
         """
         self.set(None)
 
-    def close(self):
+    def close(self) -> None:
         """
         Closes the current connection.
         This method is automatically called when the program exits using the 'with' statement.
         """
         self._send({}, OpCode.CLOSE)
-        self._socket.close()
-
-    def _connect(self):
-        pipe = self._get_pipe()
-
-        # Try to connect to a socket, starting from 0 up to 9
-        for i in range(10):
-            try:
-                self._try_socket(pipe, i)
-                break
-            except FileNotFoundError:
-                pass
-        else:
-            raise PresenceError('Cannot find a socket to connect to Discord')
-
-    def _get_pipe(self) -> str:
-        if self._platform == WINDOWS:
-            # Windows pipe
-            return R'\\.\pipe\\' + SOCKET_NAME
+        self._socket._close()
 
-        # Unix pipe
-        for env in ('XDG_RUNTIME_DIR', 'TMPDIR', 'TMP', 'TEMP'):
-            path = os.environ.get(env)
-            if path is not None:
-                return os.path.join(path, SOCKET_NAME)
-
-        return os.path.join('/tmp/', SOCKET_NAME)
-
-    def _try_socket(self, pipe: str, i: int):
-        if self._platform == WINDOWS:
-            self._socket = open(pipe.format(i), 'rb+')
-        else:
-            self._socket = socket.socket(socket.AF_UNIX)
-            self._socket.connect(pipe.format(i))
-
-    def _handshake(self):
-        data = {
-            'v': 1,
-            'client_id': self.client_id,
-        }
-        self._send(data, OpCode.HANDSHAKE)
-        _, data = self._read()
+    def _handshake(self) -> None:
+        self._send({'v': 1, 'client_id': self.client_id}, OpCode.HANDSHAKE)
+        data = self._read()
 
         if data.get('evt') != 'READY':
             raise PresenceError('Discord returned an error response after a handshake request')
 
-    def _read(self) -> tuple[int, dict[str, Any]]:
+    def _read(self) -> dict[str, Any]:
         op, length = self._read_header()
-        payload = self._read_bytes(length)
-        decoded = payload.decode('utf-8')
+        decoded = self._read_bytes(length).decode('utf-8')
         data = json.loads(decoded)
-        return op, data
+        return cast(dict[str, Any], data)
 
     def _read_header(self) -> tuple[int, int]:
-        return struct.unpack('<ii', self._read_bytes(8))
+        return cast(tuple[int, int], struct.unpack('<ii', self._read_bytes(8)))
 
     def _read_bytes(self, size: int) -> bytes:
         encoded = b''
         while size > 0:
-            if self._platform == WINDOWS:
-                encoded += self._socket.read(size)
-            else:
-                encoded += self._socket.recv(size)
-
+            encoded += self._socket._read(size)
             size -= len(encoded)
         return encoded
 
-    def _send(self, payload: dict[str, int], op: OpCode):
-        data_json = json.dumps(payload)
-        encoded = data_json.encode('utf-8')
+    def _send(self, payload: dict[str, Any], op: OpCode) -> None:
+        encoded = json.dumps(payload).encode('utf-8')
         header = struct.pack('<ii', int(op), len(encoded))
-        self._write(header + encoded)
+        self._socket._write(header + encoded)
 
-    def _write(self, data: bytes):
-        if self._platform == WINDOWS:
-            self._socket.write(data)
-            self._socket.flush()
-        else:
-            self._socket.sendall(data)
-
-    def __enter__(self):
+    def __enter__(self) -> 'Presence':
         return self
 
-    def __exit__(self, exc_type, exc_value, exc_traceback):
+    def __exit__(
+        self,
+        exc_type: Optional[type[BaseException]],
+        exc_value: Optional[BaseException],
+        exc_traceback: Optional[TracebackType],
+    ) -> None:
         self.close()
+
+class _Socket(ABC):
+    @abstractmethod
+    def __init__(self) -> None:
+        pass
+
+    @abstractmethod
+    def _read(self, size: int) -> bytes:
+        pass
+
+    @abstractmethod
+    def _write(self, data: bytes) -> None:
+        pass
+
+    @abstractmethod
+    def _close(self) -> None:
+        pass
+
+class _UnixSocket(_Socket):
+    def __init__(self) -> None:
+        pipe = os.path.join(self._get_pipe_path(), SOCKET_NAME)
+
+        # Try to connect to a socket, starting from 0 up to 9
+        for i in range(10):
+            try:
+                self._sock = socket.socket(socket.AF_UNIX)  # type: ignore [attr-defined,unused-ignore]
+                self._sock.connect(pipe.format(i))
+                break
+            except FileNotFoundError:
+                pass
+        else:
+            raise PresenceError('Cannot find a Unix socket to connect to Discord')
+
+    def _get_pipe_path(self) -> str:
+        for env in ('XDG_RUNTIME_DIR', 'TMPDIR', 'TMP', 'TEMP'):
+            path = os.environ.get(env)
+            if path is not None:
+                return path
+
+        return '/tmp/'
+
+    def _read(self, size: int) -> bytes:
+        return self._sock.recv(size)
+
+    def _write(self, data: bytes) -> None:
+        self._sock.sendall(data)
+
+    def _close(self) -> None:
+        self._sock.close()
+
+class _WindowsSocket(_Socket):
+    def __init__(self) -> None:
+        pipe = R'\\.\pipe\\' + SOCKET_NAME
+
+        # Try to connect to a socket, starting from 0 up to 9
+        for i in range(10):
+            try:
+                self._buffer = open(pipe.format(i), 'rb+')
+                break
+            except FileNotFoundError:
+                pass
+        else:
+            raise PresenceError('Cannot find a Windows socket to connect to Discord')
+
+    def _read(self, size: int) -> bytes:
+        return self._buffer.read(size)
+
+    def _write(self, data: bytes) -> None:
+        self._buffer.write(data)
+        self._buffer.flush()
+
+    def _close(self) -> None:
+        self._buffer.close()

setup.py

@@ -24,4 +24,5 @@
         'Typing :: Typed',
     ],
     python_requires='>=3.9',
+    package_data={"discordrp": ["py.typed"]},
 )