Save proxies event after interruption

Fixes #316

proxy_scraper_checker/__main__.py

@@ -114,60 +114,66 @@ async def main() -> None:
     cfg = await read_config("config.toml")
     console = Console()
     configure_logging(console=console, debug=cfg["debug"])
-
-    async with ClientSession(
-        connector=TCPConnector(ssl=http.SSL_CONTEXT),
-        headers=http.HEADERS,
-        cookie_jar=http.get_cookie_jar(),
-        raise_for_status=True,
-        fallback_charset_resolver=http.fallback_charset_resolver,
-    ) as session:
-        settings = await Settings.from_mapping(cfg, session=session)
-        storage = ProxyStorage(protocols=settings.sources)
-        with Progress(
-            TextColumn("[yellow]{task.fields[col1]}"),
-            TextColumn("[red]::"),
-            TextColumn("[green]{task.fields[col2]}"),
-            BarColumn(),
-            MofNCompleteColumn(),
-            console=console,
-            transient=True,
-        ) as progress:
-            scrape = scraper.scrape_all(
-                progress=progress,
-                session=session,
-                settings=settings,
-                storage=storage,
-            )
-            await (
-                asyncio.gather(
-                    geodb.download_geodb(progress=progress, session=session),
-                    scrape,
+    should_save = False
+    try:
+        async with ClientSession(
+            connector=TCPConnector(ssl=http.SSL_CONTEXT),
+            headers=http.HEADERS,
+            cookie_jar=http.get_cookie_jar(),
+            raise_for_status=True,
+            fallback_charset_resolver=http.fallback_charset_resolver,
+        ) as session:
+            settings = await Settings.from_mapping(cfg, session=session)
+            storage = ProxyStorage(protocols=settings.sources)
+            with Progress(
+                TextColumn("[yellow]{task.fields[col1]}"),
+                TextColumn("[red]::"),
+                TextColumn("[green]{task.fields[col2]}"),
+                BarColumn(),
+                MofNCompleteColumn(),
+                console=console,
+                transient=True,
+            ) as progress:
+                scrape = scraper.scrape_all(
+                    progress=progress,
+                    session=session,
+                    settings=settings,
+                    storage=storage,
+                )
+                await (
+                    asyncio.gather(
+                        geodb.download_geodb(
+                            progress=progress, session=session
+                        ),
+                        scrape,
+                    )
+                    if settings.enable_geolocation
+                    else scrape
+                )
+                await session.close()
+                count_before_checking = storage.get_count()
+                should_save = True
+                await checker.check_all(
+                    settings=settings,
+                    storage=storage,
+                    progress=progress,
+                    proxies_count=count_before_checking,
+                )
+    finally:
+        if should_save:
+            storage.remove_unchecked()
+            count_after_checking = storage.get_count()
+            console.print(
+                get_summary_table(
+                    before=count_before_checking, after=count_after_checking
                 )
-                if settings.enable_geolocation
-                else scrape
-            )
-            await session.close()
-            count_before_checking = storage.get_count()
-            await checker.check_all(
-                settings=settings,
-                storage=storage,
-                progress=progress,
-                proxies_count=count_before_checking,
             )
 
-    count_after_checking = storage.get_count()
-    console.print(
-        get_summary_table(
-            before=count_before_checking, after=count_after_checking
+            await output.save_proxies(storage=storage, settings=settings)
+
+        logger.info(
+            "Thank you for using https://github.com/monosans/proxy-scraper-checker"
         )
-    )
-
-    await output.save_proxies(storage=storage, settings=settings)
-
-    logger.info(
-        "Thank you for using https://github.com/monosans/proxy-scraper-checker"
-    )
 
 
 if __name__ == "__main__":

proxy_scraper_checker/output.py

@@ -53,7 +53,9 @@ def save_proxies(*, settings: Settings, storage: ProxyStorage) -> None:
                     "host": proxy.host,
                     "port": proxy.port,
                     "exit_ip": proxy.exit_ip,
-                    "timeout": round(proxy.timeout, 2),
+                    "timeout": round(proxy.timeout, 2)
+                    if proxy.timeout is not None
+                    else None,
                     "geolocation": mmdb_reader.get(proxy.exit_ip)
                     if mmdb_reader is not None and proxy.exit_ip is not None
                     else None,

proxy_scraper_checker/proxy.py

@@ -39,8 +39,12 @@ class Proxy:
     port: int
     username: str | None
     password: str | None
-    timeout: float = attrs.field(init=False, eq=False)
-    exit_ip: str | None = attrs.field(init=False, eq=False)
+    timeout: float | None = attrs.field(default=None, init=False, eq=False)
+    exit_ip: str | None = attrs.field(default=None, init=False, eq=False)
+
+    @property
+    def is_checked(self) -> bool:
+        return self.timeout is not None
 
     async def check(self, *, settings: Settings) -> None:
         async with settings.semaphore:

proxy_scraper_checker/sort.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import math
 from typing import TYPE_CHECKING
 
 from aiohttp_socks import ProxyType
@@ -19,4 +20,4 @@ def natural_sort_key(proxy: Proxy, /) -> tuple[int, ...]:
 
 
 def timeout_sort_key(proxy: Proxy, /) -> float:
-    return proxy.timeout
+    return proxy.timeout if proxy.timeout is not None else math.inf

proxy_scraper_checker/storage.py

@@ -54,5 +54,10 @@ class ProxyStorage:
             **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)