33from __future__ import annotations
44
55import os
6+ import threading
67from concurrent import futures
78from http .server import ThreadingHTTPServer
8- from typing import Any , Callable , Coroutine , Optional , TypeVar , overload
9+ from typing import Any , Callable , Coroutine , List , Optional , TypeVar , overload
910from urllib .parse import urlsplit
1011
1112from typing_extensions import ParamSpec , TypeAlias
3132 "Status" ,
3233 "all" ,
3334 "any" ,
35+ "batch" ,
3436 "call" ,
3537 "function" ,
3638 "gather" ,
4446T = TypeVar ("T" )
4547
4648_registry : Optional [Registry ] = None
47-
49+ _workers : List [Callable [None , None ]] = []
50+ _threads : List [threading .Thread ] = []
4851
4952def default_registry ():
5053 global _registry
@@ -89,6 +92,18 @@ def run(init: Optional[Callable[P, None]] = None, *args: P.args, **kwargs: P.kwa
8992 parsed_url = urlsplit ("//" + address )
9093 server_address = (parsed_url .hostname or "" , parsed_url .port or 0 )
9194 server = ThreadingHTTPServer (server_address , Dispatch (default_registry ()))
95+
96+ for worker in _workers :
97+ def entrypoint ():
98+ try :
99+ worker ()
100+ finally :
101+ server .shutdown ()
102+ _threads .append (threading .Thread (target = entrypoint ))
103+
104+ for thread in _threads :
105+ thread .start ()
106+
92107 try :
93108 if init is not None :
94109 init (* args , ** kwargs )
@@ -97,7 +112,16 @@ def run(init: Optional[Callable[P, None]] = None, *args: P.args, **kwargs: P.kwa
97112 server .shutdown ()
98113 server .server_close ()
99114
115+ for thread in _threads :
116+ thread .join ()
117+
100118
101119def batch () -> Batch :
102120 """Create a new batch object."""
103121 return default_registry ().batch ()
122+
123+
124+ def worker (fn : Callable [None , None ]) -> Callable [None , None ]:
125+ """Decorator declaring workers that will be started when dipatch.run is called."""
126+ _workers .append (fn )
127+ return fn
0 commit comments