Use AppKey (#505)

Author: Dreamsorcerer

aiohttp_debugtoolbar/main.py

@@ -1,6 +1,6 @@
 import secrets
 from pathlib import Path
-from typing import Iterable, Literal, Sequence, Tuple, Type, TypedDict, Union
+from typing import Iterable, Literal, Sequence, Type, TypedDict, Union
 
 import aiohttp_jinja2
 import jinja2
@@ -11,10 +11,12 @@
 from .panels.base import DebugPanel
 from .utils import (
     APP_KEY,
+    AppState,
     ExceptionHistory,
     STATIC_ROUTE_NAME,
     TEMPLATE_KEY,
     ToolbarStorage,
+    _Config,
 )
 from .views import ExceptionDebugView
 
@@ -35,21 +37,6 @@
 )
 
 
-class _Config(TypedDict):
-    enabled: bool
-    intercept_exc: Literal["debug", "display", False]
-    intercept_redirects: bool
-    panels: Tuple[Type[DebugPanel], ...]
-    extra_panels: Tuple[Type[DebugPanel], ...]
-    global_panels: Tuple[Type[DebugPanel], ...]
-    hosts: Sequence[str]
-    exclude_prefixes: Tuple[str, ...]
-    check_host: bool
-    button_style: str
-    max_visible_requests: int
-    path_prefix: str
-
-
 class _AppDetails(TypedDict):
     exc_history: ExceptionHistory
     pdtb_token: str
@@ -90,7 +77,6 @@ def setup(
         path_prefix=path_prefix,
     )
 
-    app[APP_KEY] = {"settings": config}
     if middleware not in app.middlewares:
         app.middlewares.append(middleware)
 
@@ -151,8 +137,13 @@ def setup(
         "GET", path_prefix, views.request_view, name="debugtoolbar.main"
     )
 
-    app[APP_KEY]["request_history"] = ToolbarStorage(max_request_history)
-    app[APP_KEY]["exc_history"] = ExceptionHistory()
-    app[APP_KEY]["pdtb_token"] = secrets.token_hex(10)
+    app[APP_KEY] = AppState(
+        {
+            "exc_history": ExceptionHistory(),
+            "pdtb_token": secrets.token_hex(10),
+            "request_history": ToolbarStorage(max_request_history),
+            "settings": config,
+        }
+    )
     if intercept_exc:
         app[APP_KEY]["exc_history"].eval_exc = intercept_exc == "debug"

aiohttp_debugtoolbar/middlewares.py

@@ -3,7 +3,7 @@
 import aiohttp_jinja2
 from aiohttp import web
 from aiohttp.typedefs import Handler
-from aiohttp.web_exceptions import _HTTPMove as HTTPMove
+from aiohttp.web_exceptions import HTTPMove
 
 from .tbtools.tbtools import get_traceback
 from .toolbar import DebugToolbar

aiohttp_debugtoolbar/utils.py

@@ -4,9 +4,15 @@
 import sys
 from collections import deque
 from itertools import islice
+from typing import Literal, Sequence, TYPE_CHECKING, Tuple, Type, TypedDict
 
-APP_KEY = "aiohttp_debugtoolbar"
-TEMPLATE_KEY = "aiohttp_debugtoolbar_jinja2"
+import jinja2
+from aiohttp.web import AppKey
+
+if TYPE_CHECKING:  # pragma: no cover
+    from .panels.base import DebugPanel
+else:
+    DebugPanel = None
 
 REDIRECT_CODES = (300, 301, 302, 303, 305, 307, 308)
 STATIC_PATH = "static/"
@@ -48,6 +54,32 @@ def __init__(self):
         self.eval_exc = "show"
 
 
+class _Config(TypedDict):
+    enabled: bool
+    intercept_exc: Literal["debug", "display", False]
+    intercept_redirects: bool
+    panels: Tuple[Type[DebugPanel], ...]
+    extra_panels: Tuple[Type[DebugPanel], ...]
+    global_panels: Tuple[Type[DebugPanel], ...]
+    hosts: Sequence[str]
+    exclude_prefixes: Tuple[str, ...]
+    check_host: bool
+    button_style: str
+    max_visible_requests: int
+    path_prefix: str
+
+
+class AppState(TypedDict):
+    exc_history: ExceptionHistory
+    pdtb_token: str
+    request_history: ToolbarStorage
+    settings: _Config
+
+
+APP_KEY = AppKey("APP_KEY", AppState)
+TEMPLATE_KEY = AppKey("TEMPLATE_KEY", jinja2.Environment)
+
+
 def addr_in(addr, hosts):
     for host in hosts:
         if ipaddress.ip_address(addr) in ipaddress.ip_network(host):

examples/extra_panels/server.py

@@ -21,12 +21,15 @@
 
 PATH_PARENT = pathlib.Path(__file__).parent
 
+db_key = web.AppKey["aiopg.Pool"]("db_key")
+redis_key = web.AppKey["aioredis.Redis"]("redis_key")
+
 
 @aiohttp_jinja2.template("index.html")
 async def basic_handler(request):
     # testing for PgSQL
     if "db" in request.app:
-        conn = await request.app["db"].acquire()
+        conn = await request.app[db_key].acquire()
         cur = await conn.cursor()
 
         await cur.execute("SELECT 1")
@@ -35,11 +38,11 @@ async def basic_handler(request):
             ret.append(row)
         assert ret == [(1,)]  # noqa: S101
 
-        await request.app["db"].release(conn)
+        await request.app[db_key].release(conn)
 
     # testing for Redis
     if "redis" in request.app:
-        with await request.app["redis"] as redis:
+        with await request.app[redis_key] as redis:
             await redis.set("TEST", "VAR", expire=5)
             assert b"VAR" == (await redis.get("TEST"))  # noqa: S101
 
@@ -55,13 +58,13 @@ async def exception_handler(request):
 
 
 async def close_pg(app):
-    app["db"].close()
-    await app["db"].wait_closed()
+    app[db_key].close()
+    await app[db_key].wait_closed()
 
 
 async def close_redis(app):
-    app["redis"].close()
-    await app["redis"].wait_closed()
+    app[redis_key].close()
+    await app[redis_key].wait_closed()
 
 
 async def init():
@@ -106,13 +109,13 @@ async def init():
         dsn = "host={host} dbname={db} user={user} password={passw} ".format(
             db="postgres", user="developer", passw="1", host="localhost"
         )
-        app["db"] = await aiopg.create_pool(dsn, minsize=1, maxsize=2)
+        app[db_key] = await aiopg.create_pool(dsn, minsize=1, maxsize=2)
         # Correct PostgreSQL shutdown
         app.on_cleanup.append(close_pg)
 
     if "aioredis" in sys.modules:
         # create redis pool
-        app["redis"] = await aioredis.Redis("127.0.0.1", 6379)
+        app[redis_key] = await aioredis.Redis()
         # Correct Redis shutdown
         app.on_cleanup.append(close_redis)
 

requirements.txt

@@ -1,7 +1,7 @@
 -e .
 
-aiohttp==3.8.5
-aiohttp-jinja2==1.5.1
+aiohttp==3.9.0
+aiohttp-jinja2==1.6
 aiohttp-mako==0.4.0
 aioredis==2.0.1
 coverage==7.2.7

setup.py

@@ -38,7 +38,7 @@ def read(fname):
     license="Apache 2",
     packages=find_packages(),
     install_requires=(
-        "aiohttp>=3.8",
+        "aiohttp>=3.9",
         "aiohttp_jinja2",
     ),
     include_package_data=True,

tests/test_panel.py

@@ -3,6 +3,7 @@
 import aiohttp_jinja2
 
 from aiohttp_debugtoolbar.panels.base import DebugPanel
+from aiohttp_debugtoolbar.utils import TEMPLATE_KEY
 
 
 async def test_request_vars_panel(create_server, aiohttp_client):
@@ -63,7 +64,7 @@ async def process_response(self, response):
     assert "pDebugToolbarHandle" in txt
 
     # check template from extra_templates
-    assert "test.jinja2" in app["aiohttp_debugtoolbar_jinja2"].list_templates()
+    assert "test.jinja2" in app[TEMPLATE_KEY].list_templates()
 
     # make sure that debug toolbar page working and extra panel exists
     resp = await client.get("/_debugtoolbar")