From 14b68a93e8c38a1e959f5dad3812234d8d250bca Mon Sep 17 00:00:00 2001 From: Vlad Pronsky Date: Thu, 18 Apr 2024 04:27:25 +0300 Subject: [PATCH] ability to raise exception when no active account (#48, #148) --- readme.md | 7 +++++-- tests/test_api.py | 25 ++++++++++++++++++++++++- twscrape/accounts_pool.py | 18 ++++++++++++++++-- twscrape/api.py | 10 +++++++--- twscrape/utils.py | 8 ++++++++ 5 files changed, 60 insertions(+), 8 deletions(-) diff --git a/readme.md b/readme.md index ae50afc..0db13ea 100644 --- a/readme.md +++ b/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. diff --git a/tests/test_api.py b/tests/test_api.py index ca44035..4229a62 100644 --- a/tests/test_api.py +++ b/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 @@ async def test_gql_params(api_mock: API, monkeypatch): 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 diff --git a/twscrape/accounts_pool.py b/twscrape/accounts_pool.py index 0ec0ac7..79114f7 100644 --- a/twscrape/accounts_pool.py +++ b/twscrape/accounts_pool.py @@ -11,7 +11,11 @@ from .account import Account 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 @@ class AccountsPool: 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 @@ class AccountsPool: msg = f'No account available for queue "{queue}". Next available at {nat}' logger.info(msg) msg_shown = True + await asyncio.sleep(5) continue else: diff --git a/twscrape/api.py b/twscrape/api.py index 72ca1f7..0a5cc29 100644 --- a/twscrape/api.py +++ b/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 diff --git a/twscrape/utils.py b/twscrape/utils.py index 7f0c57e..aa2b9d8 100644 --- a/twscrape/utils.py +++ b/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")