Skip to content

Commit

Permalink
add dispatch.worker and dispatch.batch
Browse files Browse the repository at this point in the history
Signed-off-by: Achille Roussel <achille.roussel@gmail.com>
  • Loading branch information
achille-roussel committed May 13, 2024
1 parent 833793e commit 25bd032
Showing 1 changed file with 26 additions and 2 deletions.
28 changes: 26 additions & 2 deletions src/dispatch/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,9 +3,10 @@
from __future__ import annotations

import os
import threading
from concurrent import futures
from http.server import ThreadingHTTPServer
from typing import Any, Callable, Coroutine, Optional, TypeVar, overload
from typing import Any, Callable, Coroutine, List, Optional, TypeVar, overload
from urllib.parse import urlsplit

from typing_extensions import ParamSpec, TypeAlias
Expand All @@ -31,6 +32,7 @@
"Status",
"all",
"any",
"batch",
"call",
"function",
"gather",
Expand All @@ -44,7 +46,8 @@
T = TypeVar("T")

_registry: Optional[Registry] = None

_workers: List[Callable[None, None]] = []
_threads: List[threading.Thread] = []

def default_registry():
global _registry
Expand Down Expand Up @@ -89,6 +92,18 @@ def run(init: Optional[Callable[P, None]] = None, *args: P.args, **kwargs: P.kwa
parsed_url = urlsplit("//" + address)
server_address = (parsed_url.hostname or "", parsed_url.port or 0)
server = ThreadingHTTPServer(server_address, Dispatch(default_registry()))

for worker in _workers:
def entrypoint():
try:
worker()
finally:
server.shutdown()
_threads.append(threading.Thread(target=entrypoint))

for thread in _threads:
thread.start()

try:
if init is not None:
init(*args, **kwargs)
Expand All @@ -97,7 +112,16 @@ def run(init: Optional[Callable[P, None]] = None, *args: P.args, **kwargs: P.kwa
server.shutdown()
server.server_close()

for thread in _threads:
thread.join()


def batch() -> Batch:
"""Create a new batch object."""
return default_registry().batch()


def worker(fn: Callable[None, None]) -> Callable[None, None]:
"""Decorator declaring workers that will be started when dipatch.run is called."""
_workers.append(fn)
return fn

0 comments on commit 25bd032

Please sign in to comment.