From 05d42d1ff2e050c26dd565a588eb3ff60e858305 Mon Sep 17 00:00:00 2001
From: Giovanni Barillari <g@baro.dev>
Date: Tue, 14 Mar 2023 11:16:54 +0100
Subject: [PATCH] Add `wrap_loader` param to `Server.serve` (#35)

---
 granian/server.py | 17 +++++++++++++----
 1 file changed, 13 insertions(+), 4 deletions(-)

diff --git a/granian/server.py b/granian/server.py
index b3a3ba8..c638425 100644
--- a/granian/server.py
+++ b/granian/server.py
@@ -7,7 +7,7 @@
 
 from functools import partial
 from pathlib import Path
-from typing import Any, Dict, List, Optional
+from typing import Any, Callable, Dict, List, Optional
 
 import watchfiles
 
@@ -343,17 +343,26 @@ def _serve_with_reloader(self, spawn_target, target_loader):
 
         self.shutdown()
 
-    def serve(self, spawn_target = None, target_loader = None):
+    def serve(
+        self,
+        spawn_target: Callable[..., None] = None,
+        target_loader: Callable[..., Callable[..., Any]] = None,
+        wrap_loader: bool = True
+    ):
         default_spawners = {
             Interfaces.ASGI: self._spawn_asgi_worker,
             Interfaces.RSGI: self._spawn_rsgi_worker,
             Interfaces.WSGI: self._spawn_wsgi_worker
         }
-        target_loader = target_loader or load_target
+        if target_loader:
+            if wrap_loader:
+                target_loader = partial(target_loader, self.target)
+        else:
+            target_loader = partial(load_target, self.target)
         spawn_target = spawn_target or default_spawners[self.interface]
 
         serve_method = (
             self._serve_with_reloader if self.reload_on_changes else
             self._serve
         )
-        serve_method(spawn_target, partial(target_loader, self.target))
+        serve_method(spawn_target, target_loader)