Specify explicitly the CloseCode type.

Author: aaugustin

src/websockets/asyncio/connection.py

@@ -616,7 +616,11 @@ async def send(
         else:
             raise TypeError("data must be str, bytes, iterable, or async iterable")
 
-    async def close(self, code: int = 1000, reason: str = "") -> None:
+    async def close(
+        self,
+        code: CloseCode | int = CloseCode.NORMAL_CLOSURE,
+        reason: str = "",
+    ) -> None:
         """
         Perform the closing handshake.
 

src/websockets/frames.py

@@ -360,7 +360,7 @@ class Close:
 
     """
 
-    code: int
+    code: CloseCode | int
     reason: str
 
     def __str__(self) -> str:

src/websockets/protocol.py

@@ -361,7 +361,7 @@ def send_binary(self, data: BytesLike, fin: bool = True) -> None:
         self.expect_continuation_frame = not fin
         self.send_frame(Frame(OP_BINARY, data, fin))
 
-    def send_close(self, code: int | None = None, reason: str = "") -> None:
+    def send_close(self, code: CloseCode | int | None = None, reason: str = "") -> None:
         """
         Send a `Close frame`_.
 
@@ -429,7 +429,7 @@ def send_pong(self, data: BytesLike) -> None:
             raise InvalidState(f"connection is {self.state.name.lower()}")
         self.send_frame(Frame(OP_PONG, data))
 
-    def fail(self, code: int, reason: str = "") -> None:
+    def fail(self, code: CloseCode | int, reason: str = "") -> None:
         """
         `Fail the WebSocket connection`_.
 

src/websockets/sync/connection.py

@@ -568,7 +568,11 @@ def send(
         else:
             raise TypeError("data must be str, bytes, or iterable")
 
-    def close(self, code: int = CloseCode.NORMAL_CLOSURE, reason: str = "") -> None:
+    def close(
+        self,
+        code: CloseCode | int = CloseCode.NORMAL_CLOSURE,
+        reason: str = "",
+    ) -> None:
         """
         Perform the closing handshake.