#3261

Build # #3261

Build Type

push

coveralls-python

Committed by GitHub

Commit Message

<a href="https://github.com/mosquito/aiomisc/commit/<a class=hub.com/mosquito/aiomisc/commit/ae7bcfdf2e24800292bf13854667303a1da8f498">ae7bcfdf2<a href="https://github.com/mosquito/aiomisc/commit/ae7bcfdf2e24800292bf13854667303a1da8f498">">Merge </a><a class="double-link" href="https://github.com/mosquito/aiomisc/commit/<a class="double-link" href="https://github.com/mosquito/aiomisc/commit/da4bcedcc499166867a7575455d89c3c6c5c4977">da4bcedcc</a>">da4bcedcc</a><a href="https://github.com/mosquito/aiomisc/commit/ae7bcfdf2e24800292bf13854667303a1da8f498"> into bfc83dd07">bfc83dd07</a>

Run Details

455 of 649 branches covered (70.11%)

Branch coverage included in aggregate %.

6 of 6 new or added lines in 3 files covered. (100.0%)

1812 of 2101 relevant lines covered (86.24%)

0.86 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

87.97

/aiomisc/utils.py

import asyncio
import itertools
import logging.handlers
import socket
from functools import wraps
from multiprocessing import cpu_count
from typing import (
    Any, Awaitable, Callable, Iterable, List, Optional, Tuple, TypeVar, Union,
)

from .thread_pool import ThreadPoolExecutor


T = TypeVar("T")
TimeoutType = Union[int, float]

try:
    import uvloop  # type: ignore
    event_loop_policy = uvloop.EventLoopPolicy()
except ImportError:
    event_loop_policy = asyncio.DefaultEventLoopPolicy()


log = logging.getLogger(__name__)


def chunk_list(iterable: Iterable[T], size: int) -> Iterable[List[T]]:
    """
    Split list or generator by chunks with fixed maximum size.
    """

    iterable = iter(iterable)

    item = list(itertools.islice(iterable, size))
    while item:
        yield item
        item = list(itertools.islice(iterable, size))


OptionsType = Iterable[Tuple[int, int, int]]


def bind_socket(
    *args: Any, address: str, port: int, options: OptionsType = (),
    reuse_addr: bool = True, reuse_port: bool = False,
    proto_name: Optional[str] = None
) -> socket.socket:
    """

    Bind socket and set ``setblocking(False)`` for just created socket.
    This detects ``address`` format and select socket family automatically.

    :param args: which will be passed to stdlib's socket constructor (optional)
    :param address: bind address
    :param port: bind port
    :param options: Tuple of pairs which contain socket option
                    to set and the option value.
    :param reuse_addr: set socket.SO_REUSEADDR
    :param reuse_port: set socket.SO_REUSEPORT
    :param proto_name: protocol name which will be logged after binding
    :return: socket.socket
    """

    if not args:
        if ":" in address:
            args = (socket.AF_INET6, socket.SOCK_STREAM)
        else:
            args = (socket.AF_INET, socket.SOCK_STREAM)

    sock = socket.socket(*args)
    sock.setblocking(False)

    sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, int(reuse_addr))
    if hasattr(socket, "SO_REUSEPORT"):
        sock.setsockopt(
            socket.SOL_SOCKET, socket.SO_REUSEPORT, int(reuse_port),
        )
    else:
        log.warning("SO_REUSEPORT is not implemented by underlying library.")

    for level, option, value in options:
        sock.setsockopt(level, option, value)

    unix_address_family = getattr(socket, "AF_UNIX", None)
    if sock.family == unix_address_family:
        proto_name = proto_name or "unix"
        sock.bind(address)
    else:
        proto_name = proto_name or "tcp"
        sock.bind((address, port))

    sock_addr = sock.getsockname()
    if not isinstance(sock_addr, str):
        sock_addr = sock_addr[:2]

    if sock.family == socket.AF_INET6:
        log.info("Listening %s://[%s]:%s", proto_name, *sock_addr)
    elif sock.family == unix_address_family:
        log.info("Listening %s://%s", proto_name, sock_addr)
    else:
        log.info("Listening %s://%s:%s", proto_name, *sock_addr)

    return sock


def create_default_event_loop(
    pool_size: Optional[int] = None,
    policy: asyncio.AbstractEventLoopPolicy = event_loop_policy,
    debug: bool = False,
) -> Tuple[asyncio.AbstractEventLoop, ThreadPoolExecutor]:
    """
    Creates an event loop and thread pool executor

    :param pool_size: thread pool maximal size
    :param policy: event loop policy
    :param debug: set ``loop.set_debug(True)`` if True
    """
    try:
        asyncio.get_event_loop().close()
    except RuntimeError:
        pass  # event loop is not created yet

    asyncio.set_event_loop_policy(policy)

    loop = asyncio.new_event_loop()
    loop.set_debug(debug)
    asyncio.set_event_loop(loop)

    pool_size = pool_size or cpu_count()
    thread_pool = ThreadPoolExecutor(pool_size, statistic_name="default")
    loop.set_default_executor(thread_pool)

    return loop, thread_pool


def new_event_loop(
    pool_size: Optional[int] = None,
    policy: asyncio.AbstractEventLoopPolicy = event_loop_policy,
) -> asyncio.AbstractEventLoop:
    loop, thread_pool = create_default_event_loop(pool_size, policy)
    return loop


def shield(func: Callable[..., Awaitable[T]]) -> Callable[..., Awaitable[T]]:
    """
    Simple and useful decorator for wrap the coroutine to `asyncio.shield`.

    >>> @shield
    ... async def non_cancelable_func():
    ...     await asyncio.sleep(1)

    """

    async def awaiter(future: Awaitable[Any]) -> Any:
        return await future

    @wraps(func)
    def wrap(*args: Any, **kwargs: Any) -> Awaitable[Any]:
        return wraps(func)(awaiter)(asyncio.shield(func(*args, **kwargs)))

    return wrap


class SelectResult:
    def __init__(self, length: int):
        self.length = length
        self.result_idx = None      # type: Optional[int]
        self.is_exception = None    # type: Optional[bool]
        self.value = None           # type: Any

    def set_result(self, idx: int, value: Any, is_exception: bool) -> None:
        if self.result_idx is not None:
            return

        self.value = value
        self.result_idx = idx
        self.is_exception = is_exception

    def result(self) -> T:
        if self.is_exception:
            raise self.value
        return self.value

    def done(self) -> bool:
        return self.result_idx is not None

    def __iter__(self) -> Iterable[Optional[T]]:
        for i in range(self.length):
            if i == self.result_idx:
                yield self.value
            else:
                yield None


def cancel_tasks(tasks: Iterable[asyncio.Future]) -> asyncio.Future:
    """
    All passed tasks will be cancelled and a new task will be returned.

    :param tasks: tasks which will be cancelled
    """

    future = asyncio.get_event_loop().create_future()
    future.set_result(None)

    if not tasks:
        return future

    cancelled_tasks = []
    exc = asyncio.CancelledError()

    for task in tasks:
        if task.done():
            continue

        if isinstance(task, asyncio.Task):
            task.cancel()
            cancelled_tasks.append(task)

        elif isinstance(task, asyncio.Future):
            task.set_exception(exc)

        else:
            log.warning(
                "Skipping object %r because it's not a Task or Future", task,
            )

    if not cancelled_tasks:
        return future

    waiter = asyncio.ensure_future(
        asyncio.gather(
            *cancelled_tasks, return_exceptions=True
        ),
    )

    return waiter


async def _select_waiter(
    idx: int, awaitable: Awaitable[T],
    result: SelectResult,
) -> None:
    try:
        ret = await awaitable
    except asyncio.CancelledError:
        raise
    except Exception as e:
        return result.set_result(idx, e, is_exception=True)

    result.set_result(idx, ret, is_exception=False)


async def select(
    *awaitables: Awaitable[Any],
    return_exceptions: bool = False, cancel: bool = True,
    timeout: Optional[TimeoutType] = None,
    wait: bool = True, loop: Optional[asyncio.AbstractEventLoop] = None
) -> SelectResult:
    """

    :param awaitables: awaitable objects
    :param return_exceptions: if True exception will not be raised
                              just returned as result
    :param cancel: cancel unfinished coroutines (default True)
    :param timeout: execution timeout
    :param wait: when False and ``cancel=True``, unfinished coroutines will
                 be cancelled in the background.
    :param loop: event loop
    """

    loop = loop or asyncio.get_event_loop()
    result = SelectResult(len(awaitables))

    coroutines = [
        loop.create_task(_select_waiter(idx, coroutine, result))
        for idx, coroutine in enumerate(awaitables)
    ]

    _, pending = await loop.create_task(
        asyncio.wait(
            coroutines,
            timeout=timeout,
            return_when=asyncio.FIRST_COMPLETED,
        ),
    )

    if cancel:
        cancelling = cancel_tasks(pending)

        if wait:
            await cancelling

    if result.is_exception and not return_exceptions:
        result.result()

    return result


def awaitable(
    func: Callable[..., Union[T, Awaitable[T]]],
) -> Callable[..., Awaitable[T]]:
    """

    Decorator wraps function and returns a function which returns
    awaitable object. In case than a function returns a future,
    the original future will be returned. In case then the function
    returns a coroutine, the original coroutine will be returned.
    In case than function returns non-awaitable object, it's will
    be wrapped to a new coroutine which just returns this object.
    It's useful when you don't want to check function result before
    use it in ``await`` expression.
    """

    # Avoid python 3.8+ warning
    if asyncio.iscoroutinefunction(func):
        return func     # type: ignore

    async def awaiter(obj: T) -> T:
        return obj

    @wraps(func)
    def wrap(*args: Any, **kwargs: Any) -> Awaitable[T]:
        result = func(*args, **kwargs)

        if hasattr(result, "__await__"):
            return result       # type: ignore
        if asyncio.iscoroutine(result) or asyncio.isfuture(result):
            return result       # type: ignore

        return awaiter(result)  # type: ignore

    return wrap

1	import asyncio	1✔
2	import itertools	1✔
3	import logging.handlers	1✔
4	import socket	1✔
5	from functools import wraps	1✔
6	from multiprocessing import cpu_count	1✔
7	from typing import (	1✔
8	Any, Awaitable, Callable, Iterable, List, Optional, Tuple, TypeVar, Union,
9	)
10
11	from .thread_pool import ThreadPoolExecutor	1✔
12
13
14	T = TypeVar("T")	1✔
15	TimeoutType = Union[int, float]	1✔
16
17	try:	1✔
18	import uvloop # type: ignore	1✔
19	event_loop_policy = uvloop.EventLoopPolicy()	1✔
20	except ImportError:	×
21	event_loop_policy = asyncio.DefaultEventLoopPolicy()	×
22
23
24	log = logging.getLogger(__name__)	1✔
25
26
27	def chunk_list(iterable: Iterable[T], size: int) -> Iterable[List[T]]:	1✔
28	"""
29	Split list or generator by chunks with fixed maximum size.
30	"""
31
32	iterable = iter(iterable)	1✔
33
34	item = list(itertools.islice(iterable, size))	1✔
35	while item:	1✔
36	yield item	1✔
37	item = list(itertools.islice(iterable, size))	1✔
38
39
40	OptionsType = Iterable[Tuple[int, int, int]]	1✔
41
42
43	def bind_socket(	1✔
44	*args: Any, address: str, port: int, options: OptionsType = (),
45	reuse_addr: bool = True, reuse_port: bool = False,
46	proto_name: Optional[str] = None
47	) -> socket.socket:
48	"""
49
50	Bind socket and set ``setblocking(False)`` for just created socket.
51	This detects ``address`` format and select socket family automatically.
52
53	:param args: which will be passed to stdlib's socket constructor (optional)
54	:param address: bind address
55	:param port: bind port
56	:param options: Tuple of pairs which contain socket option
57	to set and the option value.
58	:param reuse_addr: set socket.SO_REUSEADDR
59	:param reuse_port: set socket.SO_REUSEPORT
60	:param proto_name: protocol name which will be logged after binding
61	:return: socket.socket
62	"""
63
64	if not args:	1✔
65	if ":" in address:	1✔
66	args = (socket.AF_INET6, socket.SOCK_STREAM)	1✔
67	else:
68	args = (socket.AF_INET, socket.SOCK_STREAM)	1✔
69
70	sock = socket.socket(*args)	1✔
71	sock.setblocking(False)	1✔
72
73	sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, int(reuse_addr))	1✔
74	if hasattr(socket, "SO_REUSEPORT"):	1!
75	sock.setsockopt(	1✔
76	socket.SOL_SOCKET, socket.SO_REUSEPORT, int(reuse_port),
77	)
78	else:
79	log.warning("SO_REUSEPORT is not implemented by underlying library.")
80
81	for level, option, value in options:
82	sock.setsockopt(level, option, value)
83
84	unix_address_family = getattr(socket, "AF_UNIX", None)
85	if sock.family == unix_address_family:
86	proto_name = proto_name or "unix"
87	sock.bind(address)
88	else:
89	proto_name = proto_name or "tcp"
90	sock.bind((address, port))
91
92	sock_addr = sock.getsockname()
93	if not isinstance(sock_addr, str):
94	sock_addr = sock_addr[:2]
95
96	if sock.family == socket.AF_INET6:
97	log.info("Listening %s://[%s]:%s", proto_name, *sock_addr)
98	elif sock.family == unix_address_family:
99	log.info("Listening %s://%s", proto_name, sock_addr)
100	else:
101	log.info("Listening %s://%s:%s", proto_name, *sock_addr)
102
103	return sock
104
105
106	def create_default_event_loop(
107	pool_size: Optional[int] = None,
108	policy: asyncio.AbstractEventLoopPolicy = event_loop_policy,
109	debug: bool = False,
110	) -> Tuple[asyncio.AbstractEventLoop, ThreadPoolExecutor]:
111	"""
112	Creates an event loop and thread pool executor
113
114	:param pool_size: thread pool maximal size
115	:param policy: event loop policy
116	:param debug: set ``loop.set_debug(True)`` if True
117	"""
118	try:	1✔
119	asyncio.get_event_loop().close()	1✔
120	except RuntimeError:	×
121	pass # event loop is not created yet	×
122
123	asyncio.set_event_loop_policy(policy)	1✔
124
125	loop = asyncio.new_event_loop()	1✔
126	loop.set_debug(debug)	1✔
127	asyncio.set_event_loop(loop)	1✔
128
129	pool_size = pool_size or cpu_count()	1✔
130	thread_pool = ThreadPoolExecutor(pool_size, statistic_name="default")	1✔
131	loop.set_default_executor(thread_pool)	1✔
132
133	return loop, thread_pool	1✔
134
135
136	def new_event_loop(
137	pool_size: Optional[int] = None,
138	policy: asyncio.AbstractEventLoopPolicy = event_loop_policy,
139	) -> asyncio.AbstractEventLoop:
140	loop, thread_pool = create_default_event_loop(pool_size, policy)	×
141	return loop	×
142
143
144	def shield(func: Callable[..., Awaitable[T]]) -> Callable[..., Awaitable[T]]:
145	"""
146	Simple and useful decorator for wrap the coroutine to `asyncio.shield`.
147
148	>>> @shield
149	... async def non_cancelable_func():
150	... await asyncio.sleep(1)
151
152	"""
153
154	async def awaiter(future: Awaitable[Any]) -> Any:	1✔
155	return await future	1✔
156
157	@wraps(func)	1✔
158	def wrap(args: Any, *kwargs: Any) -> Awaitable[Any]:	1!
159	return wraps(func)(awaiter)(asyncio.shield(func(args, *kwargs)))	1✔
160
161	return wrap	1✔
162
163
164	class SelectResult:	1✔
165	def __init__(self, length: int):	1✔
166	self.length = length	1✔
167	self.result_idx = None # type: Optional[int]	1✔
168	self.is_exception = None # type: Optional[bool]	1✔
169	self.value = None # type: Any	1✔
170
171	def set_result(self, idx: int, value: Any, is_exception: bool) -> None:	1✔
172	if self.result_idx is not None:	1✔
173	return	1✔
174
175	self.value = value	1✔
176	self.result_idx = idx	1✔
177	self.is_exception = is_exception	1✔
178
179	def result(self) -> T:	1✔
180	if self.is_exception:	1!
181	raise self.value	1✔
182	return self.value	×
183
184	def done(self) -> bool:	1✔
185	return self.result_idx is not None	×
186
187	def __iter__(self) -> Iterable[Optional[T]]:	1✔
188	for i in range(self.length):	1✔
189	if i == self.result_idx:	1✔
190	yield self.value	1✔
191	else:
192	yield None	1✔
193
194
195	def cancel_tasks(tasks: Iterable[asyncio.Future]) -> asyncio.Future:
196	"""
197	All passed tasks will be cancelled and a new task will be returned.
198
199	:param tasks: tasks which will be cancelled
200	"""
201
202	future = asyncio.get_event_loop().create_future()	1✔
203	future.set_result(None)	1✔
204
205	if not tasks:	1✔
206	return future	1✔
207
208	cancelled_tasks = []	1✔
209	exc = asyncio.CancelledError()	1✔
210
211	for task in tasks:	1!
212	if task.done():	1✔
213	continue	1✔
214
215	if isinstance(task, asyncio.Task):	1✔
216	task.cancel()	1✔
217	cancelled_tasks.append(task)	1✔
218
219	elif isinstance(task, asyncio.Future):	1!
220	task.set_exception(exc)	1✔
221
222	else:
223	log.warning(	×
224	"Skipping object %r because it's not a Task or Future", task,
225	)
226
227	if not cancelled_tasks:
228	return future
229
230	waiter = asyncio.ensure_future(
231	asyncio.gather(
232	*cancelled_tasks, return_exceptions=True
233	),
234	)
235
236	return waiter
237
238
239	async def _select_waiter(
240	idx: int, awaitable: Awaitable[T],
241	result: SelectResult,
242	) -> None:
243	try:	1✔
244	ret = await awaitable	1✔
245	except asyncio.CancelledError:	1✔
246	raise	1✔
247	except Exception as e:	1✔
248	return result.set_result(idx, e, is_exception=True)	1✔
249
250	result.set_result(idx, ret, is_exception=False)	1✔
251
252
253	async def select(
254	*awaitables: Awaitable[Any],
255	return_exceptions: bool = False, cancel: bool = True,
256	timeout: Optional[TimeoutType] = None,
257	wait: bool = True, loop: Optional[asyncio.AbstractEventLoop] = None
258	) -> SelectResult:
259	"""
260
261	:param awaitables: awaitable objects
262	:param return_exceptions: if True exception will not be raised
263	just returned as result
264	:param cancel: cancel unfinished coroutines (default True)
265	:param timeout: execution timeout
266	:param wait: when False and ``cancel=True``, unfinished coroutines will
267	be cancelled in the background.
268	:param loop: event loop
269	"""
270
271	loop = loop or asyncio.get_event_loop()	1✔
272	result = SelectResult(len(awaitables))	1✔
273
274	coroutines = [	1!
275	loop.create_task(_select_waiter(idx, coroutine, result))
276	for idx, coroutine in enumerate(awaitables)
277	]
278
279	_, pending = await loop.create_task(
280	asyncio.wait(
281	coroutines,
282	timeout=timeout,
283	return_when=asyncio.FIRST_COMPLETED,
284	),
285	)
286
287	if cancel:
288	cancelling = cancel_tasks(pending)
289
290	if wait:
291	await cancelling
292
293	if result.is_exception and not return_exceptions:
294	result.result()
295
296	return result
297
298
299	def awaitable(	1✔
300	func: Callable[..., Union[T, Awaitable[T]]],
301	) -> Callable[..., Awaitable[T]]:
302	"""
303
304	Decorator wraps function and returns a function which returns
305	awaitable object. In case than a function returns a future,
306	the original future will be returned. In case then the function
307	returns a coroutine, the original coroutine will be returned.
308	In case than function returns non-awaitable object, it's will
309	be wrapped to a new coroutine which just returns this object.
310	It's useful when you don't want to check function result before
311	use it in ``await`` expression.
312	"""
313
314	# Avoid python 3.8+ warning
315	if asyncio.iscoroutinefunction(func):	1✔
316	return func # type: ignore	1✔
317
318	async def awaiter(obj: T) -> T:	1✔
319	return obj	1✔
320
321	@wraps(func)	1✔
322	def wrap(args: Any, *kwargs: Any) -> Awaitable[T]:	1!
323	result = func(args, *kwargs)	1✔
324
325	if hasattr(result, "__await__"):	1✔
326	return result # type: ignore	1✔
327	if asyncio.iscoroutine(result) or asyncio.isfuture(result):	1!
328	return result # type: ignore	×
329
330	return awaiter(result) # type: ignore	1✔
331		×
332	return wrap	1✔

mosquito / aiomisc / #3261

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous