Add router add_subroute method

src/turbosloth/router.py

@@ -15,10 +15,19 @@ class Route:
     regexp_subroutes: list[tuple[Pattern, list[str], Route]]
     handler: dict[MethodType, InternalHandlerType]
 
-    def __init__(self) -> None:
-        self.static_subroutes = {}
-        self.regexp_subroutes = []
-        self.handler = {}
+    def __init__(self,
+                 static_subroutes: Optional[dict[str, Route]] = None,
+                 regexp_subroutes: Optional[list[tuple[Pattern, list[str], Route]]] = None,
+                 handler: Optional[dict[MethodType, InternalHandlerType]] = None) -> None:
+        if static_subroutes is None:
+            static_subroutes = {}
+        if regexp_subroutes is None:
+            regexp_subroutes = []
+        if handler is None:
+            handler = {}
+        self.static_subroutes = static_subroutes
+        self.regexp_subroutes = regexp_subroutes
+        self.handler = handler
 
     def _find_regexp_subroute(self, p: Pattern) -> Optional[tuple[list[str], Route]]:
         for _p, _n, _r in self.regexp_subroutes:
@@ -31,10 +40,9 @@ class Route:
             self.regexp_subroutes.append((p, names, r))
             self.regexp_subroutes.sort(key=lambda x: len(x[0].pattern), reverse=True)
 
-    def add(self, method: MethodType, sequence: Sequence[str], handler: InternalHandlerType) -> None:
+    def walk_into_route(self, sequence: Sequence[str]) -> Route:
         if len(sequence) == 0:
-            self.handler[method] = handler
-            return
+            return self
 
         part = sequence[0]
 
@@ -57,8 +65,18 @@ class Route:
             if subroute is None:
                 subroute = Route()
                 self.static_subroutes[part] = subroute
+        return subroute.walk_into_route(sequence[1:])
 
-        subroute.add(method, sequence[1:], handler)
+    def add(self,
+            method: MethodType,
+            sequence: Sequence[str],
+            handler: InternalHandlerType) -> None:
+        if len(sequence) == 0:
+            self.handler[method] = handler
+            return
+
+        subroute = self.walk_into_route(sequence)
+        subroute.handler[method] = handler
 
     def get(self, method: MethodType, sequence: Sequence[str]) -> tuple[dict[str, str], Optional[InternalHandlerType]]:
         if len(sequence) == 0:
@@ -90,6 +108,32 @@ class Route:
         submatches |= matches
         return submatches, handler
 
+    def compose_routes(self, other: Route) -> Route:
+        new_static_subroutes = self.static_subroutes.copy()
+        for k, v in other.static_subroutes.items():
+            old_r = new_static_subroutes.get(k)
+            if old_r is not None:
+                v = old_r.compose_routes(v)
+            new_static_subroutes[k] = v
+
+        new_regexp_subroutes = self.regexp_subroutes.copy()
+
+        for p, n, r in other.regexp_subroutes:
+            found = False
+            for i, (_p, _n, _r) in enumerate(new_regexp_subroutes):
+                if p == _p:
+                    new_regexp_subroutes[i] = (_p, n, _r.compose_routes(r))
+                    found = True
+                    break
+            if not found:
+                new_regexp_subroutes.append((p, n, r))
+
+        return Route(
+            new_static_subroutes,
+            new_regexp_subroutes,
+            self.handler | other.handler
+        )
+
 
 class Router:
     _root: Route
@@ -132,3 +176,25 @@ class Router:
         if h is None:
             raise NotFoundException(path or '/')
         return h
+
+    def add_subroute(self, basepath: str, subr: Route | Router):
+        if isinstance(subr, Router):
+            subr = subr._root
+
+        segments = basepath.split('/')
+        while len(segments) > 0 and len(segments[0]) == 0:
+            segments = segments[1:]
+        if len(segments) > 1:
+            tgt = self._root.walk_into_route(segments[:-1])
+            old_tail = tgt.static_subroutes.get(segments[-1])
+            if old_tail is not None:
+                subr = old_tail.compose_routes(subr)
+
+            tgt.static_subroutes[segments[-1]] = subr
+        elif len(segments) == 1:
+            old_tail = self._root.static_subroutes.get(segments[-1])
+            if old_tail is not None:
+                subr = old_tail.compose_routes(subr)
+            self._root.static_subroutes[segments[0]] = subr
+        else:
+            self._root = self._root.compose_routes(subr)

tests/test_router.py

@@ -79,3 +79,32 @@ def test_router_pattern_match():
 
     with pytest.raises(NotFoundException, match='404\tNot Found: asd'):
         r.match('GET', 'asd/basdf')
+
+
+def test_subroutes():
+    async def f(*args):
+        pass
+
+    async def d(*args):
+        pass
+
+    r1 = Router()
+    r2 = Router()
+
+    r1.add('GET', '/asdf', f)
+    r2.add('GET', '/asdf/a', d)
+    r1.add_subroute('', r2)
+
+    assert r1.match('GET', '/asdf') == ({}, f)
+    assert r1.match('GET', '/asdf/a') == ({}, d)
+    r1.add_subroute('/asdf', r2)
+    assert r1.match('GET', '/asdf/asdf/a') == ({}, d)
+
+    r1.add_subroute('/asdf' * 5, r2)
+    assert r1.match('GET', '/asdf' * 5 + '/asdf/a') == ({}, d)
+
+    with pytest.raises(NotFoundException):
+        r1.match('GET', '/asdf/' * 5 + '/asdf/a')
+
+    r1.add_subroute('/asdf/' * 5, r2)
+    assert r1.match('GET', '/asdf/' * 5 + '/asdf/a') == ({}, d)