minor fix

邮件批量发送脚本.py

@@ -1,13 +1,13 @@
 import argparse
 import pandas
 import time
-import itertools
 
 from selenium import webdriver
 from selenium.common.exceptions import StaleElementReferenceException, TimeoutException
 from selenium.webdriver.common.by import By
 from selenium.webdriver.support import expected_conditions as EC
 from selenium.webdriver.support.wait import WebDriverWait
+from selenium.webdriver.remote.webelement import WebElement
 
 from tkinter import Tk, filedialog
 from datetime import datetime, timedelta
@@ -16,7 +16,7 @@ parser = argparse.ArgumentParser(description="邮件批量发送脚本")
 parser.add_argument('input', nargs='?')
 parser.add_argument('--column-address', type=str, nargs='?', default='邮箱')
 parser.add_argument('--column-code', type=str, nargs='?', default='客户编号')
-parser.add_argument('--column-mark', type=str, nargs='?', default='已发送')
+parser.add_argument('--column-sent', type=str, nargs='?', default='已发送')
 parser.add_argument('-u', '--url', type=str, nargs='?', default='https://id.ionos.fr/identifier')
 parser.add_argument('-a', '--address', type=str, required=True)
 parser.add_argument('-p', '--password', type=str, required=True)
@@ -49,7 +49,7 @@ def main():
         return 1
 
     limit = len(recipients)
-    marks = data.setdefault(args.column_mark, [None] * limit)
+    sents = data.setdefault(args.column_sent, [None] * limit)
 
     print(f'[信息] 已读取联系人信息共 {limit} 条')
     if limit == 0: return 0
@@ -63,7 +63,7 @@ def main():
         print(f'[!!!!] 初始化时发生了错误：{e}')
         return 2
 
-    def locate(selector, condition=EC.presence_of_element_located, parent=driver):
+    def locate(selector, condition=EC.presence_of_element_located, parent=driver) -> WebElement:
         while True:
             try:
                 wait = WebDriverWait(parent, timeout=args.timeout)
@@ -171,11 +171,12 @@ def main():
     while active and index < limit:
         recipient = recipients[index]
         code = codes[index]
-        mark = marks[index]
         occurrence = occurrences.setdefault(code, [0])
+
+        current = index
         index += 1
 
-        if mark is not None:
+        if sents[current] is not None:
             print(f'[信息] 已跳过项目 {recipient}')
             continue
 
@@ -211,7 +212,7 @@ def main():
                     wait = WebDriverWait(driver, timeout=args.interval)
                     alert = wait.until(lambda x: x.find_element(By.CSS_SELECTOR, "div.io-ox-alert.io-ox-alert-error"))
                 except:
-                    marks[index-1] = '✔'
+                    sents[current] = '✔'
                     occurrence[0] += 1
                     global sent
                     sent += 1
@@ -236,8 +237,7 @@ def main():
                 elif key in  ['S', 's']: break
                 else: active = False
 
-    key = input('[????] 是否需要保存到文件？确定 (Y) / 取消 (N): ')
-    if key in ['Y', 'y']:
+    if input('[????] 是否保存到文件？确定 (Y) / 取消 (N): ') in ['Y', 'y']:
         print(f'[信息] 正在写入文件：{args.input}')
         try:
             pandas.DataFrame.from_dict(data).to_excel(args.input, index=False, sheet_name='Sheet1')