Fix filename (#81)

* fix name issues

* remove magic number
This commit is contained in:
Xie Yanbo 2023-03-07 09:35:52 +08:00 committed by GitHub
parent 7412abb589
commit 2cfc89415c
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -10,7 +10,7 @@ from pathlib import Path
import openai import openai
import requests import requests
from bs4 import BeautifulSoup as bs from bs4 import BeautifulSoup as bs
from ebooklib import epub from ebooklib import epub, ITEM_DOCUMENT
from rich import print from rich import print
from tqdm import tqdm from tqdm import tqdm
@ -177,7 +177,7 @@ class BEPUB:
try: try:
for i in self.origin_book.get_items(): for i in self.origin_book.get_items():
pbar.update(index) pbar.update(index)
if i.get_type() == 9: if i.get_type() == ITEM_DOCUMENT:
soup = bs(i.content, "html.parser") soup = bs(i.content, "html.parser")
p_list = soup.findAll("p") p_list = soup.findAll("p")
is_test_done = IS_TEST and index > TEST_NUM is_test_done = IS_TEST and index > TEST_NUM
@ -198,7 +198,7 @@ class BEPUB:
break break
i.content = soup.prettify().encode() i.content = soup.prettify().encode()
new_book.add_item(i) new_book.add_item(i)
name = self.epub_name.split(".")[0] name, _ = os.path.splitext(self.epub_name)
epub.write_epub(f"{name}_bilingual.epub", new_book, {}) epub.write_epub(f"{name}_bilingual.epub", new_book, {})
pbar.close() pbar.close()
except (KeyboardInterrupt, Exception) as e: except (KeyboardInterrupt, Exception) as e:
@ -312,7 +312,7 @@ if __name__ == "__main__":
RESUME = options.resume RESUME = options.resume
if not OPENAI_API_KEY: if not OPENAI_API_KEY:
raise Exception("Need openai API key, please google how to") raise Exception("Need openai API key, please google how to")
if not options.book_name.endswith(".epub"): if not options.book_name.lower().endswith(".epub"):
raise Exception("please use epub file") raise Exception("please use epub file")
model = MODEL_DICT.get(options.model, "chatgpt") model = MODEL_DICT.get(options.model, "chatgpt")
language = options.language language = options.language