64 lines
1.7 KiB
Python
64 lines
1.7 KiB
Python
from __future__ import annotations
|
|
|
|
import itertools
|
|
from collections import Counter
|
|
from typing import TYPE_CHECKING
|
|
|
|
from . import sort
|
|
|
|
if TYPE_CHECKING:
|
|
from typing import Iterable, Iterator
|
|
|
|
from aiohttp_socks import ProxyType
|
|
|
|
from .proxy import Proxy
|
|
|
|
|
|
class ProxyStorage:
|
|
__slots__ = ("_proxies", "enabled_protocols")
|
|
|
|
def __init__(self, *, protocols: Iterable[ProxyType]) -> None:
|
|
self.enabled_protocols = set(protocols)
|
|
self._proxies: set[Proxy] = set()
|
|
|
|
def add(self, proxy: Proxy, /) -> None:
|
|
self.enabled_protocols.add(proxy.protocol)
|
|
self._proxies.add(proxy)
|
|
|
|
def remove(self, proxy: Proxy, /) -> None:
|
|
self._proxies.remove(proxy)
|
|
|
|
def get_grouped(self) -> dict[ProxyType, tuple[Proxy, ...]]:
|
|
key = sort.protocol_sort_key
|
|
return {
|
|
**{
|
|
proto: ()
|
|
for proto in sort.PROTOCOL_ORDER
|
|
if proto in self.enabled_protocols
|
|
},
|
|
**{
|
|
proto: tuple(v)
|
|
for (_, proto), v in itertools.groupby(
|
|
sorted(self, key=key), key=key
|
|
)
|
|
},
|
|
}
|
|
|
|
def get_count(self) -> dict[ProxyType, int]:
|
|
return {
|
|
**{
|
|
proto: 0
|
|
for proto in sort.PROTOCOL_ORDER
|
|
if proto in self.enabled_protocols
|
|
},
|
|
**Counter(proxy.protocol for proxy in self),
|
|
}
|
|
|
|
def remove_unchecked(self) -> None:
|
|
for p in self._proxies.copy():
|
|
if not p.is_checked:
|
|
self._proxies.remove(p)
|
|
|
|
def __iter__(self) -> Iterator[Proxy]:
|
|
return iter(self._proxies)
|