Coverage for faststream / _internal / cli / supervisors / asgi_multiprocess.py: 0%

1import inspect 

2from typing import TYPE_CHECKING, Any 

3 

4from faststream._internal._compat import HAS_UVICORN, uvicorn 

5from faststream.asgi.app import cast_uvicorn_params 

6from faststream.exceptions import INSTALL_UVICORN 

7 

8if TYPE_CHECKING: 

9 from faststream._internal.basic_types import SettingField 

10 from faststream._internal.cli.dto import RunArgs 

11 

12if HAS_UVICORN: 

13 from uvicorn.supervisors.multiprocess import Multiprocess, Process 

14 

15 class UvicornExtraConfig(uvicorn.Config): # type: ignore[misc] 

16 def __init__( 

17 self, 

18 run_extra_options: dict[str, "SettingField"], 

19 *args: Any, 

20 **kwargs: Any, 

21 ) -> None: 

22 super().__init__(*args, **kwargs) 

23 self._run_extra_options = run_extra_options 

24 

25 def load(self) -> None: 

26 super().load() 

27 self.loaded_app.app._run_extra_options = self._run_extra_options 

28 

29 class UvicornMultiprocess(Multiprocess): 

30 config: UvicornExtraConfig 

31 

32 def init_processes(self) -> None: 

33 for i in range(self.processes_num): 

34 self.config._run_extra_options["worker_id"] = i 

35 process = Process(self.config, self.target, self.sockets) 

36 process.start() 

37 self.processes.append(process) 

38 

39 

40class ASGIMultiprocess: 

41 def __init__( 

42 self, 

43 target: str, 

44 args: "RunArgs", 

45 workers: int, 

46 ) -> None: 

47 self._target = target 

48 self._run_extra_options = cast_uvicorn_params(args.extra_options) 

49 self._workers = workers 

50 self._is_factory = args.is_factory 

51 self._log_level = args.log_level 

52 

53 def run(self) -> None: 

54 if not HAS_UVICORN: 

55 raise ImportError(INSTALL_UVICORN) 

56 

57 config = UvicornExtraConfig( 

58 app=self._target, 

59 factory=self._is_factory, 

60 log_level=self._log_level, 

61 workers=self._workers, 

62 **{ 

63 key: v 

64 for key, v in self._run_extra_options.items() 

65 if key in set(inspect.signature(uvicorn.Config).parameters.keys()) 

66 }, 

67 run_extra_options=self._run_extra_options, 

68 ) 

69 server = uvicorn.Server(config) 

70 sock = config.bind_socket() 

71 UvicornMultiprocess(config, target=server.run, sockets=[sock]).run()