♻️ Replace _types module by asgiref.typing (encode#1044)

* ♻️ Replace _types module by asgiref.typing * ♻️ Use explicit hinting instead of TypeDict constructor * Update uvicorn/lifespan/on.py Co-authored-by: euri10 <[email protected]> Co-authored-by: euri10 <[email protected]>
eduflow · May 27, 2021 · 926a8f5 · 926a8f5
1 parent d28f683
commit 926a8f5
Show file tree

Hide file tree

Showing 5 changed files with 32 additions and 74 deletions.
diff --git a/setup.py b/setup.py
@@ -44,6 +44,7 @@ def get_packages(package):
 env_marker_below_38 = "python_version < '3.8'"
 
 minimal_requirements = [
+    "asgiref>=3.3.4",
     "click>=7.*",
     "h11>=0.8",
     "typing-extensions;" + env_marker_below_38,

diff --git a/uvicorn/_types.py b/uvicorn/_types.py
diff --git a/uvicorn/config.py b/uvicorn/config.py
@@ -7,7 +7,12 @@
 import socket
 import ssl
 import sys
-from typing import List, Tuple
+from typing import List, Tuple, Union
+
+if sys.version_info < (3, 8):
+    from typing_extensions import Literal
+else:
+    from typing import Literal
 
 import click
 
@@ -219,7 +224,7 @@ def __init__(
             self.forwarded_allow_ips = forwarded_allow_ips
 
     @property
-    def asgi_version(self) -> str:
+    def asgi_version(self) -> Union[Literal["2.0"], Literal["3.0"]]:
         return {"asgi2": "2.0", "asgi3": "3.0", "wsgi": "3.0"}[self.interface]
 
     @property

diff --git a/uvicorn/lifespan/on.py b/uvicorn/lifespan/on.py
@@ -1,9 +1,27 @@
 import asyncio
 import logging
 from asyncio import Queue
+from typing import Union
+
+from asgiref.typing import (
+    LifespanScope,
+    LifespanShutdownCompleteEvent,
+    LifespanShutdownEvent,
+    LifespanShutdownFailedEvent,
+    LifespanStartupCompleteEvent,
+    LifespanStartupEvent,
+    LifespanStartupFailedEvent,
+)
 
 from uvicorn import Config
-from uvicorn._types import LifespanReceiveMessage, LifespanScope, LifespanSendMessage
+
+LifespanReceiveMessage = Union[LifespanStartupEvent, LifespanShutdownEvent]
+LifespanSendMessage = Union[
+    LifespanStartupFailedEvent,
+    LifespanShutdownFailedEvent,
+    LifespanStartupCompleteEvent,
+    LifespanShutdownCompleteEvent,
+]
 
 STATE_TRANSITION_ERROR = "Got invalid state transition on lifespan protocol."
 
@@ -30,8 +48,8 @@ async def startup(self) -> None:
         main_lifespan_task = loop.create_task(self.main())  # noqa: F841
         # Keep a hard reference to prevent garbage collection
         # See https://github.com/encode/uvicorn/pull/972
-
-        await self.receive_queue.put({"type": "lifespan.startup"})
+        startup_event: LifespanStartupEvent = {"type": "lifespan.startup"}
+        await self.receive_queue.put(startup_event)
         await self.startup_event.wait()
 
         if self.startup_failed or (self.error_occured and self.config.lifespan == "on"):
@@ -44,7 +62,8 @@ async def shutdown(self) -> None:
         if self.error_occured:
             return
         self.logger.info("Waiting for application shutdown.")
-        await self.receive_queue.put({"type": "lifespan.shutdown"})
+        shutdown_event: LifespanShutdownEvent = {"type": "lifespan.shutdown"}
+        await self.receive_queue.put(shutdown_event)
         await self.shutdown_event.wait()
 
         if self.shutdown_failed or (

diff --git a/uvicorn/protocols/utils.py b/uvicorn/protocols/utils.py
@@ -2,7 +2,7 @@
 import urllib
 from typing import Optional, Tuple
 
-from uvicorn._types import WWWScope
+from asgiref.typing import WWWScope
 
 
 def get_remote_addr(transport: asyncio.Transport) -> Optional[Tuple[str, int]]: