codekingpro/portable-devtools

1"""Low level HTTP server."""
2
3import asyncio
4from typing import Any, Awaitable, Callable, Dict, List, Optional  # noqa
5
6from .abc import AbstractStreamWriter
7from .http_parser import RawRequestMessage
8from .streams import StreamReader
9from .web_protocol import RequestHandler, _RequestFactory, _RequestHandler
10from .web_request import BaseRequest
11
12__all__ = ("Server",)
13
14
15class Server:
16    def __init__(
17        self,
18        handler: _RequestHandler,
19        *,
20        request_factory: Optional[_RequestFactory] = None,
21        handler_cancellation: bool = False,
22        loop: Optional[asyncio.AbstractEventLoop] = None,
23        **kwargs: Any,
24    ) -> None:
25        self._loop = loop or asyncio.get_running_loop()
26        self._connections: Dict[RequestHandler, asyncio.Transport] = {}
27        self._kwargs = kwargs
28        # requests_count is the number of requests being processed by the server
29        # for the lifetime of the server.
30        self.requests_count = 0
31        self.request_handler = handler
32        self.request_factory = request_factory or self._make_request
33        self.handler_cancellation = handler_cancellation
34
35    @property
36    def connections(self) -> List[RequestHandler]:
37        return list(self._connections.keys())
38
39    def connection_made(
40        self, handler: RequestHandler, transport: asyncio.Transport
41    ) -> None:
42        self._connections[handler] = transport
43
44    def connection_lost(
45        self, handler: RequestHandler, exc: Optional[BaseException] = None
46    ) -> None:
47        if handler in self._connections:
48            if handler._task_handler:
49                handler._task_handler.add_done_callback(
50                    lambda f: self._connections.pop(handler, None)
51                )
52            else:
53                del self._connections[handler]
54
55    def _make_request(
56        self,
57        message: RawRequestMessage,
58        payload: StreamReader,
59        protocol: RequestHandler,
60        writer: AbstractStreamWriter,
61        task: "asyncio.Task[None]",
62    ) -> BaseRequest:
63        return BaseRequest(message, payload, protocol, writer, task, self._loop)
64
65    def pre_shutdown(self) -> None:
66        for conn in self._connections:
67            conn.close()
68
69    async def shutdown(self, timeout: Optional[float] = None) -> None:
70        coros = (conn.shutdown(timeout) for conn in self._connections)
71        await asyncio.gather(*coros)
72        self._connections.clear()
73
74    def __call__(self) -> RequestHandler:
75        try:
76            return RequestHandler(self, loop=self._loop, **self._kwargs)
77        except TypeError:
78            # Failsafe creation: remove all custom handler_args
79            kwargs = {
80                k: v
81                for k, v in self._kwargs.items()
82                if k in ["debug", "access_log_class"]
83            }
84            return RequestHandler(self, loop=self._loop, **kwargs)
85