aiomisc 17.5.17__py3-none-any.whl → 17.5.20__py3-none-any.whl

aiomisc/aggregate.py

@@ -7,8 +7,8 @@ from dataclasses import dataclass
 from inspect import Parameter
 from typing import (
     Any,
-    Awaitable,
     Callable,
+    Coroutine,
     Generic,
     Iterable,
     List,
@@ -246,7 +246,7 @@ class Aggregator(AggregatorAsync[V, R], Generic[V, R]):
 
 def aggregate(
     leeway_ms: float, max_count: Optional[int] = None
-) -> Callable[[AggregateFunc[V, R]], Callable[[V], Awaitable[R]]]:
+) -> Callable[[AggregateFunc[V, R]], Callable[[V], Coroutine[Any, Any, R]]]:
     """
     Parametric decorator that aggregates multiple
     (but no more than ``max_count`` defaulting to ``None``) single-argument
@@ -275,7 +275,9 @@ def aggregate(
 
     :return:
     """
-    def decorator(func: AggregateFunc[V, R]) -> Callable[[V], Awaitable[R]]:
+    def decorator(
+        func: AggregateFunc[V, R]
+    ) -> Callable[[V], Coroutine[Any, Any, R]]:
         aggregator = Aggregator(
             func, max_count=max_count, leeway_ms=leeway_ms,
         )
@@ -285,7 +287,10 @@ def aggregate(
 
 def aggregate_async(
     leeway_ms: float, max_count: Optional[int] = None,
-) -> Callable[[AggregateAsyncFunc[V, R]], Callable[[V], Awaitable[R]]]:
+) -> Callable[
+    [AggregateAsyncFunc[V, R]],
+    Callable[[V], Coroutine[Any, Any, R]]
+]:
     """
     Same as ``aggregate``, but with ``func`` arguments of type ``Arg``
     containing ``value`` and ``future`` attributes instead. In this setting
@@ -298,7 +303,7 @@ def aggregate_async(
     """
     def decorator(
         func: AggregateAsyncFunc[V, R]
-    ) -> Callable[[V], Awaitable[R]]:
+    ) -> Callable[[V], Coroutine[Any, Any, R]]:
         aggregator = AggregatorAsync(
             func, max_count=max_count, leeway_ms=leeway_ms,
         )

aiomisc/backoff.py

@@ -2,7 +2,7 @@ import asyncio
 import sys
 from functools import wraps
 from typing import (
-    Any, Awaitable, Callable, Optional, Tuple, Type, TypeVar, Union,
+    Any, Callable, Coroutine, Optional, Tuple, Type, TypeVar, Union,
 )
 
 from .counters import Statistic
@@ -42,7 +42,10 @@ def asyncbackoff(
     giveup: Optional[Callable[[Exception], bool]] = None,
     statistic_name: Optional[str] = None,
     statistic_class: Type[BackoffStatistic] = BackoffStatistic,
-) -> Callable[[Callable[P, Awaitable[T]]], Callable[P, Awaitable[T]]]:
+) -> Callable[
+    [Callable[P, Coroutine[Any, Any, T]]],
+    Callable[P, Coroutine[Any, Any, T]],
+]:
     """
     Patametric decorator that ensures that ``attempt_timeout`` and
     ``deadline`` time limits are met by decorated function.
@@ -85,7 +88,9 @@ def asyncbackoff(
     exceptions = tuple(exceptions) or ()
     exceptions += asyncio.TimeoutError,
 
-    def decorator(func: Callable[P, Awaitable[T]]) -> Callable[P, Awaitable[T]]:
+    def decorator(
+        func: Callable[P, Coroutine[Any, Any, T]]
+    ) -> Callable[P, Coroutine[Any, Any, T]]:
         if attempt_timeout is not None:
             func = timeout(attempt_timeout)(func)
 
@@ -145,7 +150,10 @@ def asyncretry(
     pause: Number = 0,
     giveup: Optional[Callable[[Exception], bool]] = None,
     statistic_name: Optional[str] = None,
-) -> Callable[[Callable[P, Awaitable[T]]], Callable[P, Awaitable[T]]]:
+) -> Callable[
+    [Callable[P, Coroutine[Any, Any, T]]],
+    Callable[P, Coroutine[Any, Any, T]],
+]:
     """
     Shortcut of ``asyncbackoff(None, None, 0, Exception)``.
 

aiomisc/process_pool.py

@@ -39,9 +39,7 @@ class ProcessPoolExecutor(ProcessPoolExecutorBase, EventLoopMixin):
         self._statistic.sum_time += loop.time() - start_time
 
     def submit(self, *args: Any, **kwargs: Any) -> Future:
-        """
-        Submit blocking function to the pool
-        """
+        """Submit blocking function to the pool"""
         loop = asyncio.get_running_loop()
         start_time = loop.time()
         future = super().submit(*args, **kwargs)

aiomisc/service/dns/__init__.py

@@ -0,0 +1,14 @@
+from . import records
+from .service import DNSServer, TCPDNSServer, UDPDNSServer
+from .store import DNSStore
+from .zone import DNSZone
+
+
+__all__ = (
+    "DNSServer",
+    "DNSStore",
+    "DNSZone",
+    "TCPDNSServer",
+    "UDPDNSServer",
+    "records",
+)

aiomisc/service/dns/records.py

@@ -0,0 +1,400 @@
+import enum
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from typing import Any, Hashable, Iterable, List
+
+import dnslib  # type: ignore[import-untyped]
+
+
+class RecordType(enum.IntEnum):
+    A = 1
+    A6 = 38
+    AAAA = 28
+    AFSDB = 18
+    ANY = 255
+    APL = 42
+    AXFR = 252
+    CAA = 257
+    CDNSKEY = 60
+    CDS = 59
+    CERT = 37
+    CNAME = 5
+    CSYNC = 62
+    DHCID = 49
+    DLV = 32769
+    DNAME = 39
+    DNSKEY = 48
+    DS = 43
+    EUI48 = 108
+    EUI64 = 109
+    HINFO = 13
+    HIP = 53
+    HIP_AC = 55
+    HTTPS = 65
+    IPSECKEY = 45
+    IXFR = 251
+    KEY = 25
+    KX = 36
+    LOC = 29
+    MX = 15
+    NAPTR = 35
+    NS = 2
+    NSEC = 47
+    NSEC3 = 50
+    NSEC3PARAM = 51
+    NULL = 10
+    OPENPGPKEY = 61
+    OPT = 41
+    PTR = 12
+    RP = 17
+    RRSIG = 46
+    SIG = 24
+    SOA = 6
+    SPF = 99
+    SRV = 33
+    SSHFP = 44
+    SVCB = 64
+    TA = 32768
+    TKEY = 249
+    TLSA = 52
+    TSIG = 250
+    TXT = 16
+    URI = 256
+    ZONEMD = 63
+
+
+class DNSClass(enum.IntEnum):
+    IN = 1
+    CS = 2
+    CH = 3
+    Hesiod = 4
+    NONE = 254
+    ASTERISK = 255
+
+
+class RD(dnslib.RD, Hashable, ABC):
+    def __hash__(self) -> int:
+        return hash(self.data)
+
+    @classmethod
+    @abstractmethod
+    def create(cls, *args: Any, **kwargs: Any) -> "DNSRecord":
+        raise NotImplementedError
+
+
+@dataclass(frozen=True)
+class DNSRecord:
+    name: str
+    type: RecordType
+    data: RD
+    cls: DNSClass = field(default=DNSClass.IN)
+    ttl: int = field(default=3600)
+
+    def rr(self, query_type: int) -> dnslib.RR:
+        return dnslib.RR(
+            rname=self.name,
+            rtype=query_type,
+            rclass=self.cls,
+            ttl=self.ttl,
+            rdata=self.data,
+        )
+
+
+class A(dnslib.A, RD):
+    @classmethod
+    def create(cls, name: str, ip: str, ttl: int = 3600) -> DNSRecord:
+        return DNSRecord(
+            name=name,
+            type=RecordType.A,
+            data=cls(ip),
+            ttl=ttl,
+        )
+
+
+class AAAA(dnslib.AAAA, RD):
+    @classmethod
+    def create(cls, name: str, ipv6: str, ttl: int = 3600) -> DNSRecord:
+        return DNSRecord(
+            name=name,
+            type=RecordType.AAAA,
+            data=cls(ipv6),
+            ttl=ttl,
+        )
+
+
+class RDLabel(RD):
+    label: bytes
+    __type__: RecordType
+
+    @classmethod
+    def create(cls, name: str, label: str, ttl: int = 3600) -> DNSRecord:
+        return DNSRecord(
+            name=name,
+            type=cls.__type__,
+            data=cls(label),
+            ttl=ttl,
+        )
+
+    def __hash__(self) -> int:
+        return hash(self.label)
+
+
+class CNAME(dnslib.CNAME, RDLabel):
+    __type__ = RecordType.CNAME
+
+
+class NS(dnslib.NS, RDLabel):
+    __type__ = RecordType.NS
+
+
+class PTR(dnslib.PTR, RDLabel):
+    __type__ = RecordType.PTR
+
+
+class DNAME(dnslib.DNAME, RDLabel):
+    __type__ = RecordType.DNAME
+
+
+class MX(dnslib.MX, RD):
+    @classmethod
+    def create(
+        cls, name: str, exchange: str, preference: int, ttl: int = 3600,
+    ) -> DNSRecord:
+        return DNSRecord(
+            name=name,
+            type=RecordType.MX,
+            data=cls(exchange, preference),
+            ttl=ttl,
+        )
+
+
+class TXT(dnslib.TXT, RD):
+    @classmethod
+    def create(cls, name: str, text: str, ttl: int = 3600) -> DNSRecord:
+        return DNSRecord(
+            name=name,
+            type=RecordType.TXT,
+            data=cls(text),
+            ttl=ttl,
+        )
+
+
+class SOA(dnslib.SOA, RD):
+    @classmethod
+    def create(
+        cls, name: str, mname: str, rname: str, serial: int,
+        refresh: int, retry: int, expire: int, minimum: int,
+        ttl: int = 3600,
+    ) -> DNSRecord:
+        return DNSRecord(
+            name=name,
+            type=RecordType.SOA,
+            data=cls(
+                mname, rname,
+                (serial, refresh, retry, expire, minimum),
+            ),
+            ttl=ttl,
+        )
+
+
+class SRV(dnslib.SRV, RD):
+    @classmethod
+    def create(
+        cls, name: str, priority: int, weight: int, port: int,
+        target: str, ttl: int = 3600,
+    ) -> DNSRecord:
+        return DNSRecord(
+            name=name,
+            type=RecordType.SRV,
+            data=cls(priority, weight, port, target),
+            ttl=ttl,
+        )
+
+
+class CAA(dnslib.CAA, RD):
+    @classmethod
+    def create(
+        cls, name: str, flags: int, tag: str, value: str,
+        ttl: int = 3600,
+    ) -> DNSRecord:
+        return DNSRecord(
+            name=name,
+            type=RecordType.CAA,
+            data=cls(flags, tag, value),
+            ttl=ttl,
+        )
+
+
+class NAPTR(dnslib.NAPTR, RD):
+    @classmethod
+    def create(
+        cls, name: str, order: int, preference: int, flags: str,
+        service: str, regexp: str, replacement: str, ttl: int = 3600,
+    ) -> DNSRecord:
+        return DNSRecord(
+            name=name,
+            type=RecordType.NAPTR,
+            data=cls(order, preference, flags, service, regexp, replacement),
+            ttl=ttl,
+        )
+
+
+class DS(dnslib.DS, RD):
+    @classmethod
+    def create(
+        cls, name: str, key_tag: int, algorithm: int, digest_type: int,
+        digest: str, ttl: int = 3600,
+    ) -> DNSRecord:
+        return DNSRecord(
+            name=name,
+            type=RecordType.DS,
+            data=cls(key_tag, algorithm, digest_type, digest),
+            ttl=ttl,
+        )
+
+
+class DNSKEY(dnslib.DNSKEY, RD):
+    @classmethod
+    def create(
+        cls, name: str, flags: int, protocol: int, algorithm: int,
+        key: str, ttl: int = 3600,
+    ) -> DNSRecord:
+        return DNSRecord(
+            name=name,
+            type=RecordType.DNSKEY,
+            data=cls(flags, protocol, algorithm, key),
+            ttl=ttl,
+        )
+
+
+class RRSIG(dnslib.RRSIG, RD):
+    @classmethod
+    def create(
+        cls, name: str, type_covered: int, algorithm: int, labels: int,
+        original_ttl: int, expiration: int, inception: int, key_tag: int,
+        signer: str, signature: str, ttl: int = 3600,
+    ) -> DNSRecord:
+        return DNSRecord(
+            name=name,
+            type=RecordType.RRSIG,
+            data=cls(
+                type_covered, algorithm, labels, original_ttl, expiration,
+                inception, key_tag, signer, signature,
+            ),
+            ttl=ttl,
+        )
+
+
+class NSEC(dnslib.NSEC, RD):
+    @classmethod
+    def create(
+        cls, name: str, next_domain: str,
+        rrtypes: Iterable[int], ttl: int = 3600,
+    ) -> DNSRecord:
+        return DNSRecord(
+            name=name,
+            type=RecordType.NSEC,
+            data=cls(next_domain, list(rrtypes)),
+            ttl=ttl,
+        )
+
+
+class HTTPS(dnslib.HTTPS, RD):
+    @classmethod
+    def create(
+        cls, name: str, priority: int, target: str,
+        params: List[str], ttl: int = 3600,
+    ) -> DNSRecord:
+        return DNSRecord(
+            name=name,
+            type=RecordType.HTTPS,
+            data=cls(priority, target, params),
+            ttl=ttl,
+        )
+
+
+class LOC(dnslib.LOC, RD):
+    @classmethod
+    def create(
+        cls, name: str, latitude: float, longitude: float,
+        altitude: float, size: float, h_precision: float,
+        v_precision: float, ttl: int = 3600,
+    ) -> DNSRecord:
+        return DNSRecord(
+            name=name,
+            type=RecordType.LOC,
+            data=cls(
+                latitude, longitude, altitude, size, h_precision, v_precision,
+            ),
+            ttl=ttl,
+        )
+
+
+class RP(dnslib.RP, RD):
+    @classmethod
+    def create(
+        cls, name: str, mbox: str, txt: str, ttl: int = 3600,
+    ) -> DNSRecord:
+        return DNSRecord(
+            name=name,
+            type=RecordType.RP,
+            data=cls(mbox, txt),
+            ttl=ttl,
+        )
+
+
+class TLSA(dnslib.TLSA, RD):
+    @classmethod
+    def create(
+        cls, name: str, usage: int, selector: int, mtype: int, cert: str,
+        ttl: int = 3600,
+    ) -> DNSRecord:
+        return DNSRecord(
+            name=name,
+            type=RecordType.TLSA,
+            data=cls(usage, selector, mtype, cert),
+            ttl=ttl,
+        )
+
+
+class SSHFP(dnslib.SSHFP, RD):
+    @classmethod
+    def create(
+        cls, name: str, algorithm: int, fptype: int,
+        fingerprint: str, ttl: int = 3600,
+    ) -> DNSRecord:
+        return DNSRecord(
+            name=name,
+            type=RecordType.SSHFP,
+            data=cls(algorithm, fptype, fingerprint),
+            ttl=ttl,
+        )
+
+
+__all__ = (
+    "AAAA",
+    "CAA",
+    "CNAME",
+    "DNAME",
+    "DNSKEY",
+    "DS",
+    "HTTPS",
+    "LOC",
+    "MX",
+    "NAPTR",
+    "NS",
+    "NSEC",
+    "PTR",
+    "RD",
+    "RP",
+    "RRSIG",
+    "SOA",
+    "SRV",
+    "SSHFP",
+    "TLSA",
+    "TXT",
+    "A",
+    "DNSRecord",
+    "DNSClass",
+    "RecordType",
+)

aiomisc/service/dns/service.py

@@ -0,0 +1,131 @@
+import asyncio
+import logging
+from asyncio import StreamReader, StreamWriter
+from struct import Struct
+from typing import Optional, Tuple
+
+import dnslib  # type: ignore[import-untyped]
+
+from aiomisc.service import TCPServer, UDPServer
+
+from .records import RecordType
+from .store import DNSStore
+
+
+log = logging.getLogger(__name__)
+
+
+class DNSServer:
+    store: DNSStore
+
+    __proto__: str
+
+    @staticmethod
+    def get_edns_max_size(record: dnslib.DNSRecord) -> Optional[int]:
+        for rr in record.ar:
+            if rr.rtype != dnslib.QTYPE.OPT:
+                continue
+            return rr.edns_len
+        return None
+
+    def handle_request(
+        self, addr: tuple, data: bytes, should_truncate: bool = False,
+    ) -> Tuple[bytes, tuple]:
+        record = dnslib.DNSRecord.parse(data)
+        question: dnslib.DNSQuestion = record.get_q()
+        reply = record.reply()
+        query_name = str(question.get_qname())
+        query_type = question.qtype
+
+        try:
+            log.debug(
+                "Processing %s request from %r:\n%s\n",
+                self.__proto__, addr, record,
+            )
+            edns_max_size = self.get_edns_max_size(record)
+            if edns_max_size is not None:
+                reply.add_ar(dnslib.EDNS0(udp_len=edns_max_size))
+
+            records = self.store.query(query_name, RecordType(query_type))
+            for rec in records:
+                reply.add_answer(rec.rr(query_type))
+
+            if not records:
+                reply.header.rcode = dnslib.RCODE.NXDOMAIN
+
+            log.debug(
+                "Sending %s answer to %r:\n%s\n",
+                self.__proto__, addr, reply,
+            )
+
+            reply_body = reply.pack()
+
+            if len(reply_body) > 512:
+                reply.header.tc = 1
+                reply_body = reply.pack()
+
+            if should_truncate:
+                reply_body = reply_body[:edns_max_size]
+        except Exception as e:
+            log.exception(
+                "Failed to process %s request from %r: %s",
+                self.__proto__, addr, e,
+            )
+            reply = record.reply()
+            reply.header.set_rcode(dnslib.RCODE.SERVFAIL)
+            reply_body = reply.pack()
+
+        return reply_body, addr
+
+
+class UDPDNSServer(DNSServer, UDPServer):
+    MAX_SIZE = 512
+
+    store: DNSStore
+
+    __required__ = ("store",)
+    __proto__ = "UDP"
+
+    async def handle_datagram(self, data: bytes, addr: tuple) -> None:
+        self.sendto(
+            *self.handle_request(
+                data=data, addr=addr, should_truncate=True,
+            ),
+        )
+
+
+TCP_HEADER_STRUCT: Struct = Struct("!H")
+
+
+class TCPDNSServer(DNSServer, TCPServer):
+    store: DNSStore
+
+    __required__ = ("store",)
+
+    __proto__ = "TCP"
+
+    async def handle_client(
+        self, reader: StreamReader, writer: StreamWriter,
+    ) -> None:
+        addr = writer.get_extra_info("peername")
+        try:
+            while True:
+                data = await reader.readexactly(TCP_HEADER_STRUCT.size)
+                if not data:
+                    break
+                length = TCP_HEADER_STRUCT.unpack(data)[0]
+                packet = await reader.read(length)
+
+                if not data:
+                    break
+
+                reply_body, _ = self.handle_request(data=packet, addr=addr)
+
+                writer.write(TCP_HEADER_STRUCT.pack(len(reply_body)))
+                writer.write(reply_body)
+                await writer.drain()
+        except asyncio.IncompleteReadError:
+            pass
+        finally:
+            writer.close()
+            await writer.wait_closed()

aiomisc/service/dns/store.py

@@ -0,0 +1,51 @@
+from typing import Optional, Sequence, Tuple
+
+from .records import DNSRecord, RecordType
+from .tree import RadixTree
+from .zone import DNSZone
+
+
+class DNSStore:
+    zones: RadixTree[DNSZone]
+
+    __slots__ = ("zones",)
+
+    def __init__(self) -> None:
+        self.zones = RadixTree()
+
+    def add_zone(self, zone: DNSZone) -> None:
+        zone_tuple = self.get_reverse_tuple(zone.name)
+        if self.zones.search(zone_tuple):
+            raise ValueError(f"Zone {zone.name} already exists.")
+        self.zones.insert(zone_tuple, zone)
+
+    def remove_zone(self, zone_name: str) -> None:
+        zone_tuple = self.get_reverse_tuple(zone_name)
+        if not self.zones.search(zone_tuple):
+            raise ValueError(
+                f"Zone {zone_name} does not exist.",
+            )
+        # Clear zone from RadixTree
+        self.zones.insert(zone_tuple, None)
+
+    def get_zone(self, zone_name: str) -> Optional[DNSZone]:
+        zone_tuple = self.get_reverse_tuple(zone_name)
+        return self.zones.search(zone_tuple)
+
+    def query(self, name: str, record_type: RecordType) -> Sequence[DNSRecord]:
+        if not name.endswith("."):
+            name += "."
+        zone_tuple = self.get_zone_for_name(name)
+        if not zone_tuple:
+            return ()
+        zone = self.zones.search(zone_tuple)
+        return zone.get_records(name, record_type) if zone is not None else ()
+
+    def get_zone_for_name(self, name: str) -> Optional[Tuple[str, ...]]:
+        labels = self.get_reverse_tuple(name)
+        result = self.zones.find_prefix(labels)
+        return result[0] if result else None
+
+    @staticmethod
+    def get_reverse_tuple(zone_name: str) -> Tuple[str, ...]:
+        return tuple(zone_name.strip(".").split("."))[::-1]

aiomisc/service/dns/tree.py

@@ -0,0 +1,59 @@
+from typing import Any, Dict, Generic, Hashable, List, Optional, Tuple, TypeVar
+
+
+K = Tuple[str, ...]
+T = TypeVar("T", bound=Any)
+
+
+class RadixNode(Generic[T]):
+    __slots__ = ("children", "value")
+
+    def __init__(self) -> None:
+        self.children: Dict[Hashable, RadixNode[T]] = {}
+        self.value: Optional[T] = None
+
+
+class RadixTree(Generic[T]):
+    root: RadixNode[T]
+
+    __slots__ = ("root",)
+
+    def __init__(self) -> None:
+        self.root = RadixNode()
+
+    def insert(self, key: K, value: Optional[T]) -> None:
+        node = self.root
+        for part in key:
+            if part not in node.children:
+                node.children[part] = RadixNode()
+            node = node.children[part]
+        node.value = value
+
+    def search(self, key: K) -> Optional[T]:
+        node = self.root
+        for part in key:
+            if part not in node.children:
+                return None
+            node = node.children[part]
+        return node.value
+
+    def find_prefix(self, key: K) -> Optional[Tuple[K, T]]:
+        node = self.root
+        longest_prefix: List[str] = []
+        value = None
+        part: Hashable
+        if node.value is not None:
+            value = node.value
+        for part in key:
+            if part in node.children:
+                node = node.children[part]
+                longest_prefix.append(part)
+                if node.value is not None:
+                    value = node.value
+            else:
+                break
+
+        if value is None:
+            return None
+
+        return tuple(longest_prefix), value

aiomisc/service/dns/zone.py

@@ -0,0 +1,44 @@
+from collections import defaultdict
+from typing import DefaultDict, Sequence, Set, Tuple
+
+from .records import DNSRecord, RecordType
+
+
+class DNSZone:
+    records: DefaultDict[Tuple[str, RecordType], Set[DNSRecord]]
+    name: str
+
+    __slots__ = ("name", "records")
+
+    def __init__(self, name: str):
+        if not name.endswith("."):
+            name += "."
+        self.name = name
+        self.records = defaultdict(set)
+
+    def add_record(self, record: DNSRecord) -> None:
+        if not self._is_valid_record(record):
+            raise ValueError(
+                f"Record {record.name} does not belong to zone {self.name}",
+            )
+        key = (record.name, record.type)
+        self.records[key].add(record)
+
+    def remove_record(self, record: DNSRecord) -> None:
+        key = (record.name, record.type)
+        if key in self.records:
+            self.records[key].discard(record)
+            if self.records[key]:
+                return
+            del self.records[key]
+
+    def get_records(
+        self, name: str, record_type: RecordType,
+    ) -> Sequence[DNSRecord]:
+        if not name.endswith("."):
+            name += "."
+        key = (name, record_type)
+        return tuple(self.records.get(key, ()))
+
+    def _is_valid_record(self, record: DNSRecord) -> bool:
+        return record.name.endswith(self.name)

aiomisc/service/grpc_server.py

@@ -2,9 +2,12 @@ import asyncio
 import logging
 import re
 import sys
+from collections import defaultdict
 from concurrent.futures import Executor
 from types import MappingProxyType
-from typing import Any, Optional, Sequence, Set, Tuple
+from typing import (
+    Any, DefaultDict, Dict, Mapping, Optional, Sequence, Set, Tuple,
+)
 
 from .base import Service
 
@@ -36,6 +39,7 @@ class GRPCService(Service):
     _server_args: MappingProxyType
     _insecure_ports: Set[Tuple[str, PortFuture]]
     _secure_ports: Set[Tuple[str, grpc.ServerCredentials, PortFuture]]
+    _registered_services: DefaultDict[str, Dict[str, grpc.RpcMethodHandler]]
 
     def __init__(
         self, *,
@@ -60,6 +64,7 @@ class GRPCService(Service):
         self._insecure_ports = set()
         self._secure_ports = set()
         self._reflection = reflection
+        self._registered_services = defaultdict(dict)
         super().__init__(**kwds)
 
     @classmethod
@@ -90,6 +95,12 @@ class GRPCService(Service):
             service_names.append(reflection.SERVICE_NAME)
             reflection.enable_server_reflection(service_names, self._server)
 
+        for name, handlers in self._registered_services.items():
+            # noinspection PyUnresolvedReferences
+            self._server.add_registered_method_handlers(   # type: ignore
+                name, handlers,
+            )
+
         self._server.add_generic_rpc_handlers(tuple(self._services))
         await self._server.start()
 
@@ -102,6 +113,11 @@ class GRPCService(Service):
         for service in generic_rpc_handlers:
             self._services.add(service)
 
+    def add_registered_method_handlers(
+        self, name: str, handlers: Mapping[str, grpc.RpcMethodHandler],
+    ) -> None:
+        self._registered_services[name].update(handlers)
+
     def add_insecure_port(self, address: str) -> PortFuture:
         future: PortFuture = asyncio.Future()
         self._insecure_ports.add((address, future))

aiomisc/thread_pool.py

@@ -224,9 +224,7 @@ class ThreadPoolExecutor(ThreadPoolExecutorBase):
     def submit(  # type: ignore
         self, fn: F, *args: Any, **kwargs: Any,
     ) -> asyncio.Future:
-        """
-        Submit blocking function to the pool
-        """
+        """Submit blocking function to the pool"""
         if fn is None or not callable(fn):
             raise ValueError("First argument must be callable")
 

aiomisc/timeout.py

@@ -1,7 +1,7 @@
 import asyncio
 import sys
 from functools import wraps
-from typing import Awaitable, Callable, TypeVar, Union
+from typing import Any, Callable, Coroutine, TypeVar, Union
 
 
 if sys.version_info >= (3, 10):
@@ -17,10 +17,13 @@ Number = Union[int, float]
 
 def timeout(
     value: Number
-) -> Callable[[Callable[P, Awaitable[T]]], Callable[P, Awaitable[T]]]:
+) -> Callable[
+    [Callable[P, Coroutine[Any, Any, T]]],
+    Callable[P, Coroutine[Any, Any, T]],
+]:
     def decorator(
-        func: Callable[P, Awaitable[T]],
-    ) -> Callable[P, Awaitable[T]]:
+        func: Callable[P, Coroutine[Any, Any, T]],
+    ) -> Callable[P, Coroutine[Any, Any, T]]:
         if not asyncio.iscoroutinefunction(func):
             raise TypeError("Function is not a coroutine function")
 

aiomisc/version.py

@@ -2,5 +2,5 @@
 # BY: poem-plugins "git" plugin
 # NEVER EDIT THIS FILE MANUALLY
 
-version_info = (17, 5, 17)
-__version__ = "17.5.17"
+version_info = (17, 5, 20)
+__version__ = "17.5.20"

{aiomisc-17.5.17.dist-info → aiomisc-17.5.20.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: aiomisc
-Version: 17.5.17
+Version: 17.5.20
 Summary: aiomisc - miscellaneous utils for asyncio
 Home-page: https://github.com/aiokitchen/aiomisc
 License: MIT
@@ -44,6 +44,7 @@ Provides-Extra: aiohttp
 Provides-Extra: asgi
 Provides-Extra: carbon
 Provides-Extra: cron
+Provides-Extra: dns
 Provides-Extra: grpc
 Provides-Extra: raven
 Provides-Extra: rich
@@ -55,6 +56,7 @@ Requires-Dist: aiohttp-asgi (>=0.5.2,<0.6.0) ; extra == "asgi"
 Requires-Dist: asgiref (>=3.7,<4.0) ; extra == "uvicorn"
 Requires-Dist: colorlog (>=6.0,<7.0)
 Requires-Dist: croniter (==2.0) ; extra == "cron"
+Requires-Dist: dnslib (>=0.9,<0.10) ; extra == "dns"
 Requires-Dist: grpcio (>=1.56,<2.0) ; extra == "grpc"
 Requires-Dist: grpcio-reflection (>=1.56,<2.0) ; extra == "grpc"
 Requires-Dist: grpcio-tools (>=1.56,<2.0) ; extra == "grpc"

{aiomisc-17.5.17.dist-info → aiomisc-17.5.20.dist-info}/RECORD

@@ -1,7 +1,7 @@
 aiomisc/__init__.py,sha256=mgLaoGB3-WTAnrUfEN9nM_0Z_XU6xojkx1ISmW204UA,2253
 aiomisc/_context_vars.py,sha256=28A7j_NABitKMtpxuFvxQ2wCr-Fq79dAC3YjqVLLeTQ,689
-aiomisc/aggregate.py,sha256=HU_kVLzA-BdySOH3UIfTEuB2Ty1uUC200hZf4upI504,8834
-aiomisc/backoff.py,sha256=BgdT_MEByFJPyEHjfghC7WhnEen2Lpf4p2VLZV_KVi0,5605
+aiomisc/aggregate.py,sha256=4yAsJhtRc-ikm9hXer05DgYhJFGf94IXeWJRTCSSDpw,8898
+aiomisc/backoff.py,sha256=v8G4A-z9o6ou4oa-hSv3UtLPX2Km8PLsEGqHHAFtu24,5701
 aiomisc/circuit_breaker.py,sha256=nZVLuGg4rKxn84CHaIvRfBYumgdwx2VZloF8OprQKuk,12581
 aiomisc/compat.py,sha256=aYVe0J-2CvAzUHPAULNhrfMLFYNKN-z3Sg7nlBkzfxw,3291
 aiomisc/context.py,sha256=j2YRNGDAJbrg94OkFyIMyYKHKzHRRL8tSAWma1yxNKE,1549
@@ -15,7 +15,7 @@ aiomisc/periodic.py,sha256=OFYZbPkcGgeMjBk8zwsLr2TqPRTS6MNewaL3q9qK5js,2252
 aiomisc/plugins/__init__.py,sha256=eHKGec_217rBjXGf8u-Joyf7dzpO1O0QAuFan1IEyYE,1057
 aiomisc/plugins/__main__.py,sha256=y3mykRQmimDRHb_PvY4n08vegjWTjd9CONFbCecAxxw,1365
 aiomisc/pool.py,sha256=kmTEziX7U1kXxmNsTEain9Jkzn8ZbFubuCYWaozZoRk,5764
-aiomisc/process_pool.py,sha256=tYSWEXxCP-QXum2TFhyQzyvg86JUSo7EUcrCI7si2-c,1614
+aiomisc/process_pool.py,sha256=EwzB_HqxoKX7h3A7DVeRy1AIg-rN0nSpCl9KfGW54Ss,1596
 aiomisc/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 aiomisc/recurring.py,sha256=LJ13P_OhZM6mVOed8iUrS-M9PXY12516XBZcFtO7XtI,4923
 aiomisc/service/__init__.py,sha256=_JSkI4Q78UljG4CvhK6He2UGLxJaQS6cVwZ9QT0cSWo,607
@@ -24,7 +24,13 @@ aiomisc/service/asgi.py,sha256=S6ejgg7GpHfUzuj46g8QuLXTHn0YWwqI_pAGbXfoNjE,1583
 aiomisc/service/base.py,sha256=6xvD27Sh2jy-4JWer6n31FuSq798DZhLI730YYGQbYc,4322
 aiomisc/service/carbon.py,sha256=o7x2LYUOfCcWYDAP9EFdCTtsduXhqkyMAuIqVio9YsM,1743
 aiomisc/service/cron.py,sha256=VM-FjdAhy_zFkU-YTdpBSxF79upvrsCo-3kdI9cZTOo,1975
-aiomisc/service/grpc_server.py,sha256=zfUItQlxAMj2ugb9ekKX_5Pxb_VCHOv5Q4U-iOKpx9g,3928
+aiomisc/service/dns/__init__.py,sha256=RT8_0H8LrQe8sz2xvA2BvLi6CiJVZeVFltNRAx9MN7o,254
+aiomisc/service/dns/records.py,sha256=dN_yS3KKmRCF4fNk-fxRB2m-NCQaho42thG5OLFpyZU,8756
+aiomisc/service/dns/service.py,sha256=F8KSmFry6WpA4tWNrKcXVhS6GWb_Cyl3uanGu4lbOPw,3616
+aiomisc/service/dns/store.py,sha256=FchfLSoWAI2x4n5RyIcnx_i7RQjHmBbco4L4rSAanGU,1749
+aiomisc/service/dns/tree.py,sha256=q9VRcVxggrClUHyGZZqqSByMaHdU4nYPC1NdT17cVWM,1575
+aiomisc/service/dns/zone.py,sha256=Dc3bL7L3rz38J0_R790d_uJlisfJ78vwfdsxmtVfg_w,1343
+aiomisc/service/grpc_server.py,sha256=G1iRGX3jAfv3fpzNnw_n6UuhORmIZK4cz6aAi7t7qcU,4553
 aiomisc/service/periodic.py,sha256=BEKGWxWDOX11sfo7MFYp-hiB5iZ4-c2BRsjS56k5B-0,1241
 aiomisc/service/process.py,sha256=mZf8muZJNWQo95pUmzchMk08MJMhLifFK8m4CnNX49k,4477
 aiomisc/service/profiler.py,sha256=6KiJsU7tD5eO4YKZXYujV2E8P-G8GUqLAGIl0AuPNG4,1503
@@ -36,10 +42,10 @@ aiomisc/service/tracer.py,sha256=_dxk5y2JEteki9J1OXnOkI-EowD9vakSfsLaRDB4uMQ,282
 aiomisc/service/udp.py,sha256=_uHzMAkpd7y-yt3tE9jN2llEETG7r47g2DF1SO7xyig,3616
 aiomisc/service/uvicorn.py,sha256=dR-d5ge3KsCIEJfJ6SiFQ7rIyX8rKWNuGV2ZrONToj4,4323
 aiomisc/signal.py,sha256=_iiC2jukXg7-LLirIl1YATlKIIsKLbmTNFr1Ezheu7g,1728
-aiomisc/thread_pool.py,sha256=Wx0LskSv1dGWoen1lEFRacjVco62hHn6oDaM_XKH6uo,14035
-aiomisc/timeout.py,sha256=rTipPBz0GOqDZG2euVxU_6PjI63nFD_bDsy_DOHHBDQ,863
+aiomisc/thread_pool.py,sha256=591u5HV1aBmBRcaVm4kAtMLPZOb6veqhT9wkts_knqE,14017
+aiomisc/timeout.py,sha256=OhXCvQGbPZA9IQ7cRT7ZlZnsbmDgH5ioErQyjeqIEzY,919
 aiomisc/utils.py,sha256=6yTfTpeRCVzfZp-MJCmB1oayOHUBVwQwzC3U5PBAjn8,11919
-aiomisc/version.py,sha256=_IbaaVDHszvzzXVbR3WZs2Lfbpv3_HZvwbcOahEVYIc,156
+aiomisc/version.py,sha256=gJ8ARCLt56OmyswE8_WbmF4gyB3IG7xwweORF-0BZu8,156
 aiomisc/worker_pool.py,sha256=GA91KdOrBlqHthbVSTxu_d6BsBIbl-uKqW2NxqSafG0,11107
 aiomisc_log/__init__.py,sha256=ZD-Q-YTWoZdxJpMqMNMIraA_gaN0QawgnVpXz6mxd2o,4855
 aiomisc_log/enum.py,sha256=_zfCZPYCGyI9KL6TqHiYVlOfA5U5MCbsuCuDKxDHdxg,1549
@@ -57,8 +63,8 @@ aiomisc_worker/process_inner.py,sha256=8ZtjCSLrgySW57OIbuGrpEWxfysRLYKx1or1YaAqx
 aiomisc_worker/protocol.py,sha256=1smmlBbdreSmnrxuhHaUMUC10FO9xMIEcedhweQJX_A,2705
 aiomisc_worker/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 aiomisc_worker/worker.py,sha256=f8nCFhlKh84UUBUaEgCllwMRvVZiD8_UUXaeit6g3T8,3236
-aiomisc-17.5.17.dist-info/COPYING,sha256=Ky_8CQMaIixfyOreUBsl0hKN6A5fLnPF8KPQ9molMYA,1125
-aiomisc-17.5.17.dist-info/METADATA,sha256=lckQD0C-qe2_vlVonPgPr8wHb_XNdS_nP7-mVJwN8Lg,15664
-aiomisc-17.5.17.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-aiomisc-17.5.17.dist-info/entry_points.txt,sha256=KRsSPCwKJyGTWrvzpwbS0yIDwzsgDA2X6f0CBWYmNao,55
-aiomisc-17.5.17.dist-info/RECORD,,
+aiomisc-17.5.20.dist-info/COPYING,sha256=Ky_8CQMaIixfyOreUBsl0hKN6A5fLnPF8KPQ9molMYA,1125
+aiomisc-17.5.20.dist-info/METADATA,sha256=ytR2ply0GE1ehivdvQItBTdGr5PgoxeQkCmcDll4xfY,15737
+aiomisc-17.5.20.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+aiomisc-17.5.20.dist-info/entry_points.txt,sha256=KRsSPCwKJyGTWrvzpwbS0yIDwzsgDA2X6f0CBWYmNao,55
+aiomisc-17.5.20.dist-info/RECORD,,