From 8acfe41bba3bed696cab05c012da524abf857eec Mon Sep 17 00:00:00 2001
From: Vlad Pronsky <v.pronsky@gmail.com>
Date: Sat, 29 Jun 2024 16:21:42 +0300
Subject: [PATCH] fix lint

---
 pyproject.toml   |  2 +-
 twscrape/imap.py | 11 +++++++++--
 2 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index be49c14..62b0bc3 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -28,7 +28,7 @@ dependencies = [
 
 [project.optional-dependencies]
 dev = [
-  "pyright>=1.1.359",
+  "pyright>=1.1.369",
   "pytest-asyncio>=0.23.3",
   "pytest-cov>=4.1.0",
   "pytest-httpx>=0.28.0",
diff --git a/twscrape/imap.py b/twscrape/imap.py
index c615861..acf4917 100644
--- a/twscrape/imap.py
+++ b/twscrape/imap.py
@@ -7,8 +7,15 @@ from datetime import datetime
 
 from .logger import logger
 
-TWS_WAIT_EMAIL_CODE = [os.getenv("TWS_WAIT_EMAIL_CODE"), os.getenv("LOGIN_CODE_TIMEOUT"), 30]
-TWS_WAIT_EMAIL_CODE = [int(x) for x in TWS_WAIT_EMAIL_CODE if x is not None][0]
+
+def env_int(key: str | list[str], default: int) -> int:
+    key = [key] if isinstance(key, str) else key
+    val = [os.getenv(k) for k in key]
+    val = [int(x) for x in val if x is not None]
+    return val[0] if val else default
+
+
+TWS_WAIT_EMAIL_CODE = env_int(["TWS_WAIT_EMAIL_CODE", "LOGIN_CODE_TIMEOUT"], 30)
 
 
 class EmailLoginError(Exception):