优化

src/coop.py

@@ -221,6 +221,7 @@ def main():
 if __name__ == '__main__':
     try:
         country = helper.get_input_with_default("国家", "usa")
+        domain.domain_page(page, country)
         main()
     except KeyboardInterrupt:
         pass

src/helper/domain.py

 # coding: utf-8
+import time
+
+from lxml import etree
+
 
 def switch_domain(country):
     domain = "https://vendorcentral.amazon.com/"
@@ -15,3 +19,21 @@ def switch_domain(country):
         domain = "https://vendorcentral.amazon.com.mx/"
 
     return domain
+
+
+def domain_page(page,country):
+    url = switch_domain(country)
+    page.get(url)
+    while True:
+        time.sleep(1)
+        title = page.title
+        if title == "Amazon Sign-In":
+            print("请登录账号")
+            continue
+
+        tree = etree.HTML(page.html)
+        element = tree.xpath('//*[@id="cvf-submit-otp-button-announce"]')
+        if element:
+            print('请输入验证码')
+        else:
+            break
\ No newline at end of file

src/invoices.py

@@ -329,6 +329,7 @@ def main():
 if __name__ == '__main__':
     try:
         country = helper.get_input_with_default("国家", "usa")
+        domain.domain_page(page, country)
         main()
     except KeyboardInterrupt:
         pass

src/returns.py

@@ -6,6 +6,7 @@ import time
 import pandas as pd
 from DrissionPage import ChromiumPage
 from DrissionPage.errors import ElementNotFoundError
+from lxml import etree
 
 from helper import helper, excel, file, domain
 
@@ -22,23 +23,12 @@ helper.make_dir(download_path)
 # 设置下载路径，确保在打开浏览器前设置
 page.set.download_path(download_path)
 
+
 def page_get(url):
     host = domain.switch_domain(country)
     full_url = host + url
     page.get(full_url, timeout=3)
 
-    while True:
-        time.sleep(1)
-        title = page.title
-        if title == "Amazon Sign-In":
-            print("请登录账号")
-            continue
-        try:
-            page.ele("#cvf-submit-otp-button-announce")
-            print("请输入验证码")
-        except ElementNotFoundError:
-            break
-
 
 def sku_relations():
     relations_dict = {}
@@ -132,6 +122,7 @@ def main():
 if __name__ == '__main__':
     try:
         country = helper.get_input_with_default("国家", "usa")
+        domain.domain_page(page, country)
         main()
     except KeyboardInterrupt:
         pass