增加了一页仅两张图片生成文档的功能

main.py

@@ -9,35 +9,30 @@ import fitz
 
 from docx2pdf import convert
 
+from sys import exit
+
 import os
 
-def done():
+def done(exit_code = 0):
     print("正在移除多余生成文件...")
     for gen in gened_list:
         os.remove(gen)
     print("移除完毕！")
-    input("按任意键退出...")
-    exit()
-
-    
-def deal_file_loss(path, dir_name):
-    print(f" {dir_name} 文件夹中文件个数不符，请检查！")
-    done()
+    input("按回车键退出...")
+    exit(exit_code)
 
 def convert_all(path = os.path.abspath('.')):
     print("正在扫描并处理文件中...")
     items = os.listdir(path)
     for item in items:
         if os.path.isdir(item):
-            if len(os.listdir(path + '\\' + item)) < 3:
-                deal_file_loss(path, item)
             result = convert_img(item, path + '\\' + item)
             if result == 1:
                 print("PDF 转换失败！")
-                done()
+                done(1)
             elif result == 2:
                 print("JPG/JPEG 转换失败！")
-                done()
+                done(1)
     print("转换完毕！")
 
 """传入绝对路径"""
@@ -89,16 +84,18 @@ def gen_filelist(path = os.path.abspath('.')):
     for dir in dir_list:
         file_list[dir] = os.listdir(path + '\\' + dir)
     for dir, dir_file in file_list.items():
-        png_count = 0
         dir_file_copy = copy(dir_file)
         for file in dir_file_copy:
-            if ".png" in file.lower():
-                png_count += 1
-            else:
+            if ".png" not in file.lower():
                 dir_file.remove(file)
-        if png_count != 3:
+        if len(dir_file) != 3:
             print(f"在{dir}文件夹发现错误：png文件个数不符")
-            done()
+            for file in dir_file:
+                print('- ' + file)
+            if len(dir_file) == 2:          # 说明有两个文件
+                option = input("文件夹仅有两个文件，仍继续生成？（y/N）").lower()
+                if option != 'y':
+                    done(1)                     
     print("创建完毕！将要加入文档的文件如下：")
     for dir, dir_file in file_list.items():
         print('- ' + dir)
@@ -110,32 +107,31 @@ def gen_filelist(path = os.path.abspath('.')):
     return file_list
 
 def gen_docx(path = os.path.abspath('.')):
-    while True:
-        doc = Document()
-        file_list = gen_filelist()
-        input("按回车键确认...")
-        print("生成 .docx 文档中...")
-        if len(file_list) != 0:
-            for dir, dir_file in file_list.items():
-                parent_path = path + '\\' + dir + '\\'
-                for file in dir_file:
-                    if file[0:3] == "Pdf":
-                        dir_file.remove(file)
-                        dir_file.insert(0, file)
-                doc.add_picture(parent_path + dir_file[0], height = Inches(2.5))
-                table = doc.add_table(rows = 1, cols = 2)
-                cell1 = table.cell(0, 0)
-                cell1.paragraphs[0].add_run().add_picture(parent_path + dir_file[1], height = Inches(5.0))
-                cell2 = table.cell(0, 1)
+    doc = Document()
+    file_list = gen_filelist()
+    input("按回车键确认...")
+    print("生成 .docx 文档中...")
+    if len(file_list) != 0:
+        for dir, dir_file in file_list.items():
+            parent_path = path + '\\' + dir + '\\'
+            for file in dir_file:
+                if file[0:3] == "Pdf":
+                    dir_file.remove(file)
+                    dir_file.insert(0, file)                            # 将发票文件放在第一个位置
+            doc.add_picture(parent_path + dir_file[0], height = Inches(2.5))
+            table = doc.add_table(rows = 1, cols = 2)
+            cell1 = table.cell(0, 0)
+            cell1.paragraphs[0].add_run().add_picture(parent_path + dir_file[1], height = Inches(5.0))
+            cell2 = table.cell(0, 1)
+            if len(dir_file) == 3:
                 cell2.paragraphs[0].add_run().add_picture(parent_path + dir_file[2], height = Inches(5.0))
-                if (dir, dir_file) != list(file_list.items())[-1]:          # dict.items() 返回值需先转换为列表，才能索引
-                    doc.add_page_break()
-        else:
-            print("请检查文件命名是否正确！")
-            continue
-        doc.save("output.docx")
-        print("生成完毕！")
-        break
+            if (dir, dir_file) != list(file_list.items())[-1]:          # dict.items() 返回值需先转换为列表，才能索引
+                doc.add_page_break()
+    else:
+        print("请检查文件结构是否正确！")
+        done(1)
+    doc.save("output.docx")
+    print("生成完毕！")
     
 def pdf2img(pdf_path, pdf_name, zoom_x = 3, zoom_y = 3):
     doc = fitz.open(pdf_path + "\\" + pdf_name)  # 打开文档
@@ -159,7 +155,7 @@ if __name__ == "__main__":
     print("生成 .pdf 文件中...")
     convert("output.docx", "output.pdf")
     print("生成完毕！")
-    done()
+    done(0)