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

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,7 +107,6 @@ 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("按回车键确认...")
@@ -121,21 +117,21 @@ def gen_docx(path = os.path.abspath('.')):
             for file in dir_file:
                 if file[0:3] == "Pdf":
                     dir_file.remove(file)
-                        dir_file.insert(0, 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
+        print("请检查文件结构是否正确！")
+        done(1)
     doc.save("output.docx")
     print("生成完毕！")
-        break
     
 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)