From ed6cb1b8d0b3a6121ecef06fcec8e3f4de417a5b Mon Sep 17 00:00:00 2001
From: Aymeric Augustin <aymeric.augustin@fractalideas.com>
Date: Sat, 21 Oct 2023 15:26:22 +0200
Subject: [PATCH] Make it explicit that status codes can be int.

The code and tests already support it but the types didn't reflect it.

Refs #1406.
---
 docs/reference/types.rst        | 2 ++
 src/websockets/__init__.py      | 3 +++
 src/websockets/exceptions.py    | 3 ++-
 src/websockets/legacy/server.py | 8 ++++----
 src/websockets/server.py        | 3 ++-
 src/websockets/typing.py        | 7 +++++++
 6 files changed, 20 insertions(+), 6 deletions(-)

diff --git a/docs/reference/types.rst b/docs/reference/types.rst
index 88550d08d..9d3aa8310 100644
--- a/docs/reference/types.rst
+++ b/docs/reference/types.rst
@@ -7,6 +7,8 @@ Types
 
     .. autodata:: LoggerLike
 
+    .. autodata:: StatusLike
+
     .. autodata:: Origin
 
     .. autodata:: Subprotocol
diff --git a/src/websockets/__init__.py b/src/websockets/__init__.py
index 2b9c3bc54..fdb028f4c 100644
--- a/src/websockets/__init__.py
+++ b/src/websockets/__init__.py
@@ -61,6 +61,7 @@
     "ExtensionName",
     "ExtensionParameter",
     "LoggerLike",
+    "StatusLike",
     "Origin",
     "Subprotocol",
 ]
@@ -115,6 +116,7 @@
         ExtensionParameter,
         LoggerLike,
         Origin,
+        StatusLike,
         Subprotocol,
     )
 else:
@@ -176,6 +178,7 @@
             "ExtensionParameter": ".typing",
             "LoggerLike": ".typing",
             "Origin": ".typing",
+            "StatusLike": "typing",
             "Subprotocol": ".typing",
         },
         deprecated_aliases={
diff --git a/src/websockets/exceptions.py b/src/websockets/exceptions.py
index 0f0686872..f7169e3b1 100644
--- a/src/websockets/exceptions.py
+++ b/src/websockets/exceptions.py
@@ -34,6 +34,7 @@
 from typing import Optional
 
 from . import datastructures, frames, http11
+from .typing import StatusLike
 
 
 __all__ = [
@@ -330,7 +331,7 @@ class AbortHandshake(InvalidHandshake):
 
     def __init__(
         self,
-        status: http.HTTPStatus,
+        status: StatusLike,
         headers: datastructures.HeadersLike,
         body: bytes = b"",
     ) -> None:
diff --git a/src/websockets/legacy/server.py b/src/websockets/legacy/server.py
index c9b32c417..7c24dd74a 100644
--- a/src/websockets/legacy/server.py
+++ b/src/websockets/legacy/server.py
@@ -45,7 +45,7 @@
 )
 from ..http import USER_AGENT
 from ..protocol import State
-from ..typing import ExtensionHeader, LoggerLike, Origin, Subprotocol
+from ..typing import ExtensionHeader, LoggerLike, Origin, StatusLike, Subprotocol
 from .compatibility import asyncio_timeout
 from .handshake import build_response, check_request
 from .http import read_request
@@ -57,7 +57,7 @@
 
 HeadersLikeOrCallable = Union[HeadersLike, Callable[[str, Headers], HeadersLike]]
 
-HTTPResponse = Tuple[http.HTTPStatus, HeadersLike, bytes]
+HTTPResponse = Tuple[StatusLike, HeadersLike, bytes]
 
 
 class WebSocketServerProtocol(WebSocketCommonProtocol):
@@ -349,7 +349,7 @@ async def process_request(
             request_headers: request headers.
 
         Returns:
-            Optional[Tuple[http.HTTPStatus, HeadersLike, bytes]]: :obj:`None`
+            Optional[Tuple[StatusLike, HeadersLike, bytes]]: :obj:`None`
             to continue the WebSocket handshake normally.
 
             An HTTP response, represented by a 3-uple of the response status,
@@ -943,7 +943,7 @@ class Serve:
             It defaults to ``"Python/x.y.z websockets/X.Y"``.
             Setting it to :obj:`None` removes the header.
         process_request (Optional[Callable[[str, Headers], \
-            Awaitable[Optional[Tuple[http.HTTPStatus, HeadersLike, bytes]]]]]):
+            Awaitable[Optional[Tuple[StatusLike, HeadersLike, bytes]]]]]):
             Intercept HTTP request before the opening handshake.
             See :meth:`~WebSocketServerProtocol.process_request` for details.
         select_subprotocol: Select a subprotocol supported by the client.
diff --git a/src/websockets/server.py b/src/websockets/server.py
index 872e4e5af..191660553 100644
--- a/src/websockets/server.py
+++ b/src/websockets/server.py
@@ -32,6 +32,7 @@
     ExtensionHeader,
     LoggerLike,
     Origin,
+    StatusLike,
     Subprotocol,
     UpgradeProtocol,
 )
@@ -480,7 +481,7 @@ def select_subprotocol(protocol, subprotocols):
 
     def reject(
         self,
-        status: http.HTTPStatus,
+        status: StatusLike,
         text: str,
     ) -> Response:
         """
diff --git a/src/websockets/typing.py b/src/websockets/typing.py
index e672ba006..cc3e3ec0d 100644
--- a/src/websockets/typing.py
+++ b/src/websockets/typing.py
@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import http
 import logging
 from typing import List, NewType, Optional, Tuple, Union
 
@@ -7,6 +8,7 @@
 __all__ = [
     "Data",
     "LoggerLike",
+    "StatusLike",
     "Origin",
     "Subprotocol",
     "ExtensionName",
@@ -30,6 +32,11 @@
 """Types accepted where a :class:`~logging.Logger` is expected."""
 
 
+StatusLike = Union[http.HTTPStatus, int]
+"""
+Types accepted where an :class:`~http.HTTPStatus` is expected."""
+
+
 Origin = NewType("Origin", str)
 """Value of a ``Origin`` header."""