add a bunch of typehints

[hublib] / hub / static.py

diff --git a/hub/static.py b/hub/static.py
index 8c50f77915fef08adfaa8e8536c3bc13a81de71c..f3e0d844ad68d6459801f36ba3abd6f8580ff1c7 100644 (file)
--- a/hub/static.py
+++ b/hub/static.py
@@ -1,6 +1,8 @@
 from pathlib import Path
 
 from falcon.constants import MEDIA_HTML, MEDIA_JS, MEDIA_TEXT
+from falcon.request import Request
+from falcon.response import Response
 from falcon.status_codes import HTTP_OK
 from jinja2 import Template
 
@@ -13,29 +15,28 @@ class StaticFile:
     """
     Basic static file wrapper.
     """
-    media_types_per_suffix = {
+    media_types_per_suffix: dict[str, str] = {
         ".html": MEDIA_HTML,
         ".js": MEDIA_JS,
         ".css": MEDIA_CSS,
         ".txt": MEDIA_TEXT,
     }
-    content: str | None
 
-    def __init__(self, path):
+    def __init__(self, path: Path):
         self.path = path
         self.name = path.name
         self.media_type = self.get_media_type(path)
-        self.mtime = None
-        self.content = None
+        self.mtime: float | None = None
+        self.content: bytes | None = None
 
     @classmethod
-    def get_media_type(cls, path):
+    def get_media_type(cls, path: Path):
         return cls.media_types_per_suffix[path.suffix]
 
-    def get(self):
+    def get(self) -> bytes:
         mtime = self.path.stat().st_mtime
         if mtime != self.mtime:
-            with open(self.path) as fh:
+            with open(self.path, "rb") as fh:
                 self.content = fh.read()
             self.mtime = mtime
         return self.content
@@ -102,17 +103,17 @@ class TreeFileApp:
         self.base_uri = base_uri.rstrip("/")
         self.name = name or self.create_name(self.root.base_uri, self.base_uri)
         assert self.name, self.name
-        self.files_per_uris = {}
+        self.files_per_uris: dict[str, StaticFile] = {}
         for path in self.scan_files(base_dir):
             static_file = self.get_file(path)
             uri = self.uri(static_file.path)
             self.files_per_uris[uri] = static_file
             root.add_route(uri, self)
 
-    async def on_get(self, req, resp):
+    async def on_get(self, req: Request, resp: Response):
         resource = self.files_per_uris[req.path]
         resp.content_type = resource.media_type
-        resp.data = resource.get().encode()
+        resp.data = resource.get()
         resp.status = HTTP_OK
 
 
@@ -120,13 +121,13 @@ class StaticTemplateFile(StaticFile):
     TEMPLATE_SUFFIX = ".j2"
     content: Template | None
 
-    def __init__(self, path, hubapp):
+    def __init__(self, path: Path, hubapp: "TemplateTreeFileApp"):
         super().__init__(path)
         self.name = path.stem
         self.hubapp = hubapp
         self.context = {"static_file": self}
 
-    def get(self) -> str:
+    def get(self) -> bytes:
         mtime = self.path.stat().st_mtime
         if mtime != self.mtime:
             env = self.hubapp.root.env
@@ -134,10 +135,10 @@ class StaticTemplateFile(StaticFile):
                 str(self.path.relative_to(env.loader.searchpath[0]))
             )
             self.mtime = mtime
-        return self.content.render(self.context)
+        return self.content.render(self.context).encode()
 
     @classmethod
-    def get_media_type(cls, path):
+    def get_media_type(cls, path: Path) -> str:
         assert path.suffix == cls.TEMPLATE_SUFFIX
         return cls.media_types_per_suffix[path.suffixes[-2]]
 
@@ -160,9 +161,7 @@ class HubApp(TemplateTreeFileApp):
         super().__init__(app, base_dir, base_uri)
         self.ws_app = WebSocketApp(self)
         for suffix in ("client", "master"):
-            uri = self.uri(Path(f"ws_{suffix}"))
-            self.files_per_uris[uri] = self.ws_app
-            app.add_route(uri, self.ws_app, suffix=suffix)
+            app.add_route(self.uri(Path(f"ws_{suffix}")), self.ws_app, suffix=suffix)
 
     @staticmethod
     def is_master_uri(path: Path) -> bool: