ability to raise exception when no active account (#48, #148)

readme.md

@@ -183,8 +183,6 @@ twscrape login_accounts
 
 `twscrape` will start login flow for each new account. If X will ask to verify email and you provided `email_password` in `add_account`, then `twscrape` will try to receive verification code by IMAP protocol. After success login account cookies will be saved to db file for future use.
 
-_Note:_ You can increase timeout for verification code with `TWS_WAIT_EMAIL_CODE` environment variable (default: `40`, in seconds).
-
 #### Manual email verification
 
 In case your email provider not support IMAP protocol (ProtonMail, Tutanota, etc) or IMAP is disabled in settings, you can enter email verification code manually. To do this run login command with `--manual` flag.
@@ -306,6 +304,11 @@ So if you want to use proxy PER ACCOUNT, do NOT override proxy with env variable
 
 _Note:_ If proxy not working, exception will be raised from API class.
 
+## Environment variables
+
+- `TWS_WAIT_EMAIL_CODE` – timeout for email verification code during login (default: `30`, in seconds)
+- `TWS_RAISE_WHEN_NO_ACCOUNT` – raise `NoAccountError` exception when no available accounts right now, instead of waiting for availability (default: `false`, possible value: `false` / `0` / `true` / `1`)
+
 ## Limitations
 
 After 1 July 2023 Twitter [introduced new limits](https://twitter.com/elonmusk/status/1675187969420828672) and still continue to update it periodically.

tests/test_api.py

@@ -1,6 +1,11 @@
+import os
+
+import pytest
+
+from twscrape.accounts_pool import NoAccountError
 from twscrape.api import API
 from twscrape.logger import set_log_level
-from twscrape.utils import gather
+from twscrape.utils import gather, get_env_bool
 
 set_log_level("DEBUG")
 
@@ -40,3 +45,21 @@ def mock_gql_items(*a, **kw):
         assert len(args) == 1, f"{func} not called once"
         assert args[0][1]["limit"] == 100, f"limit not changed in {func}"
         assert args[0][0][1]["count"] == 100, f"count not changed in {func}"
+
+
+async def test_raise_when_no_account(api_mock: API):
+    await api_mock.pool.delete_accounts(["user1"])
+    assert len(await api_mock.pool.get_all()) == 0
+
+    assert get_env_bool("TWS_RAISE_WHEN_NO_ACCOUNT") is False
+    os.environ["TWS_RAISE_WHEN_NO_ACCOUNT"] = "1"
+    assert get_env_bool("TWS_RAISE_WHEN_NO_ACCOUNT") is True
+
+    with pytest.raises(NoAccountError):
+        await gather(api_mock.search("foo", limit=10))
+
+    with pytest.raises(NoAccountError):
+        await api_mock.user_by_id(123)
+
+    del os.environ["TWS_RAISE_WHEN_NO_ACCOUNT"]
+    assert get_env_bool("TWS_RAISE_WHEN_NO_ACCOUNT") is False

twscrape/accounts_pool.py

@@ -11,7 +11,11 @@
 from .db import execute, fetchall, fetchone
 from .logger import logger
 from .login import LoginConfig, login
-from .utils import parse_cookies, utc
+from .utils import get_env_bool, parse_cookies, utc
+
+
+class NoAccountError(Exception):
+    pass
 
 
 class AccountInfo(TypedDict):
@@ -32,9 +36,15 @@ class AccountsPool:
     # _order_by: str = "RANDOM()"
     _order_by: str = "username"
 
-    def __init__(self, db_file="accounts.db", login_config: LoginConfig | None = None):
+    def __init__(
+        self,
+        db_file="accounts.db",
+        login_config: LoginConfig | None = None,
+        raise_when_no_account=False,
+    ):
         self._db_file = db_file
         self._login_config = login_config or LoginConfig()
+        self._raise_when_no_account = raise_when_no_account
 
     async def load_from_file(self, filepath: str, line_format: str):
         line_delim = guess_delim(line_format)
@@ -270,6 +280,9 @@ async def get_for_queue_or_wait(self, queue: str) -> Account | None:
         while True:
             account = await self.get_for_queue(queue)
             if not account:
+                if self._raise_when_no_account or get_env_bool("TWS_RAISE_WHEN_NO_ACCOUNT"):
+                    raise NoAccountError(f"No account available for queue {queue}")
+
                 if not msg_shown:
                     nat = await self.next_available_at(queue)
                     if not nat:
@@ -279,6 +292,7 @@ async def get_for_queue_or_wait(self, queue: str) -> Account | None:
                     msg = f'No account available for queue "{queue}". Next available at {nat}'
                     logger.info(msg)
                     msg_shown = True
+
                 await asyncio.sleep(5)
                 continue
             else:

twscrape/api.py

@@ -59,14 +59,18 @@ class API:
     pool: AccountsPool
 
     def __init__(
-        self, pool: AccountsPool | str | None = None, debug=False, proxy: str | None = None
+        self,
+        pool: AccountsPool | str | None = None,
+        debug=False,
+        proxy: str | None = None,
+        raise_when_no_account=False,
     ):
         if isinstance(pool, AccountsPool):
             self.pool = pool
         elif isinstance(pool, str):
-            self.pool = AccountsPool(pool)
+            self.pool = AccountsPool(db_file=pool, raise_when_no_account=raise_when_no_account)
         else:
-            self.pool = AccountsPool()
+            self.pool = AccountsPool(raise_when_no_account=raise_when_no_account)
 
         self.proxy = proxy
         self.debug = debug

twscrape/utils.py

@@ -1,5 +1,6 @@
 import base64
 import json
+import os
 from collections import defaultdict
 from datetime import datetime, timezone
 from typing import Any, AsyncGenerator, Callable, TypeVar
@@ -206,3 +207,10 @@ def parse_cookies(val: str) -> dict[str, str]:
         pass
 
     raise ValueError(f"Invalid cookie value: {val}")
+
+
+def get_env_bool(key: str, default_val: bool = False) -> bool:
+    val = os.getenv(key)
+    if val is None:
+        return default_val
+    return val.lower() in ("1", "true", "yes")