25개 이상의 토픽을 선택하실 수 없습니다. Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 
 
 

125 lines
4.1 KiB

  1. from dataclasses import dataclass
  2. from typing import Any, Callable, Generic, List, Mapping, Optional, Sequence, TypeVar
  3. from ..abc import (
  4. ByteReceiveStream, ByteSendStream, ByteStream, Listener, ObjectReceiveStream, ObjectSendStream,
  5. ObjectStream, TaskGroup)
  6. T_Item = TypeVar('T_Item')
  7. T_Stream = TypeVar('T_Stream')
  8. @dataclass(eq=False)
  9. class StapledByteStream(ByteStream):
  10. """
  11. Combines two byte streams into a single, bidirectional byte stream.
  12. Extra attributes will be provided from both streams, with the receive stream providing the
  13. values in case of a conflict.
  14. :param ByteSendStream send_stream: the sending byte stream
  15. :param ByteReceiveStream receive_stream: the receiving byte stream
  16. """
  17. send_stream: ByteSendStream
  18. receive_stream: ByteReceiveStream
  19. async def receive(self, max_bytes: int = 65536) -> bytes:
  20. return await self.receive_stream.receive(max_bytes)
  21. async def send(self, item: bytes) -> None:
  22. await self.send_stream.send(item)
  23. async def send_eof(self) -> None:
  24. await self.send_stream.aclose()
  25. async def aclose(self) -> None:
  26. await self.send_stream.aclose()
  27. await self.receive_stream.aclose()
  28. @property
  29. def extra_attributes(self) -> Mapping[Any, Callable[[], Any]]:
  30. return {**self.send_stream.extra_attributes, **self.receive_stream.extra_attributes}
  31. @dataclass(eq=False)
  32. class StapledObjectStream(Generic[T_Item], ObjectStream[T_Item]):
  33. """
  34. Combines two object streams into a single, bidirectional object stream.
  35. Extra attributes will be provided from both streams, with the receive stream providing the
  36. values in case of a conflict.
  37. :param ObjectSendStream send_stream: the sending object stream
  38. :param ObjectReceiveStream receive_stream: the receiving object stream
  39. """
  40. send_stream: ObjectSendStream[T_Item]
  41. receive_stream: ObjectReceiveStream[T_Item]
  42. async def receive(self) -> T_Item:
  43. return await self.receive_stream.receive()
  44. async def send(self, item: T_Item) -> None:
  45. await self.send_stream.send(item)
  46. async def send_eof(self) -> None:
  47. await self.send_stream.aclose()
  48. async def aclose(self) -> None:
  49. await self.send_stream.aclose()
  50. await self.receive_stream.aclose()
  51. @property
  52. def extra_attributes(self) -> Mapping[Any, Callable[[], Any]]:
  53. return {**self.send_stream.extra_attributes, **self.receive_stream.extra_attributes}
  54. @dataclass(eq=False)
  55. class MultiListener(Generic[T_Stream], Listener[T_Stream]):
  56. """
  57. Combines multiple listeners into one, serving connections from all of them at once.
  58. Any MultiListeners in the given collection of listeners will have their listeners moved into
  59. this one.
  60. Extra attributes are provided from each listener, with each successive listener overriding any
  61. conflicting attributes from the previous one.
  62. :param listeners: listeners to serve
  63. :type listeners: Sequence[Listener[T_Stream]]
  64. """
  65. listeners: Sequence[Listener[T_Stream]]
  66. def __post_init__(self) -> None:
  67. listeners: List[Listener[T_Stream]] = []
  68. for listener in self.listeners:
  69. if isinstance(listener, MultiListener):
  70. listeners.extend(listener.listeners)
  71. del listener.listeners[:] # type: ignore[attr-defined]
  72. else:
  73. listeners.append(listener)
  74. self.listeners = listeners
  75. async def serve(self, handler: Callable[[T_Stream], Any],
  76. task_group: Optional[TaskGroup] = None) -> None:
  77. from .. import create_task_group
  78. async with create_task_group() as tg:
  79. for listener in self.listeners:
  80. tg.start_soon(listener.serve, handler, task_group)
  81. async def aclose(self) -> None:
  82. for listener in self.listeners:
  83. await listener.aclose()
  84. @property
  85. def extra_attributes(self) -> Mapping[Any, Callable[[], Any]]:
  86. attributes: dict = {}
  87. for listener in self.listeners:
  88. attributes.update(listener.extra_attributes)
  89. return attributes