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 requests
from bs4 import BeautifulSoup as bs
from ebooklib import epub
from ebooklib import epub, ITEM_DOCUMENT
from rich import print
from tqdm import tqdm
@ -177,7 +177,7 @@ class BEPUB:
try:
for i in self.origin_book.get_items():
pbar.update(index)
if i.get_type() == 9:
if i.get_type() == ITEM_DOCUMENT:
soup = bs(i.content, "html.parser")
p_list = soup.findAll("p")
is_test_done = IS_TEST and index > TEST_NUM
@ -198,7 +198,7 @@ class BEPUB:
break
i.content = soup.prettify().encode()
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, {})
pbar.close()
except (KeyboardInterrupt, Exception) as e:
@ -312,7 +312,7 @@ if __name__ == "__main__":
RESUME = options.resume
if not OPENAI_API_KEY:
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")
model = MODEL_DICT.get(options.model, "chatgpt")
language = options.language