利用Python给PDF添加目录

# -*- coding: UTF-8 -*-
# add.py
# by devxzh@qq.com or onexzh@gmail.com
import fitz
import re
import sys


def print_help():
    print("""
          Command is as follow:
          >>> python add.py test.pdf list.txt num
          where  'test.pdf' is the target file.
                 'list.txt' is the catalog file to be added.
                 'num' is an integer containing +/- ,the default is 0.
          """)


argc = len(sys.argv)  # Number of input parameters

if argc < 3:
    print_help()
    exit()
else:
    pdfName = sys.argv[1]  # pdf file name
    txtName = sys.argv[2]  # txt file name
    # Offset the deviation between the PDF page number and the actual page number
    offset = 0 if argc == 3 else int(sys.argv[3])

if '.pdf' in pdfName:
    try:
        doc = fitz.open(pdfName)
    except:
        print("Please check the pdf file name and try again!")
        exit()
else:
    print("source file must be .pdf !")
    exit()

toc = []  # doc.getToC()  # get contents


def countdot(str):
    """
    没有点是一级标题，1个点是二级标题，2个点是三级标题
    """
    n = 1
    for i in str:
        if i == '.':
            n += 1
    return n


def get_page(str):
    """
    :param str:
    :return: int page
    """
    page = re.findall(r"[0-9]{1,3}$", str)  # list
    return int(page[0])


def get_title(str):
    """
    :param str:
    :return: str title
    """
    return re.sub(r"[0-9]{1,3}$", '', str)


if '.txt' in txtName:
    try:
        contentFile = open(txtName, encoding='UTF-8', errors='ignore')
    except:
        print("Please check the txt file name and try again!")
        exit()

    lines = contentFile.readlines()  # get all lines

    for line in lines:  # get every line
        # Skip blank lines and comment lines
        if line[0] == '#' or len(line.split()) == 0:
            continue
        part = line.split()  # Split line into three elements
        # if len(part) != 3:  # 缺少一项或多出一项(多出空格导致)
        #     print("Error: {0} lose some element".format(part))
        #     exit()
        print(part)
        if '章' in part[0]:
            level = 1
        else:
            level = 2
        page = get_page(part[1]) + int(offset)
        title = get_title(part[1])

        # print("{0},{1},{2}".format(level, part[1], page))
        toc.append([level, title, page])

    doc.setToC(toc)  # set contents
    newName = eval(repr(pdfName).replace("\\", "-"))  # Remove the '\' in the path
    partName = newName.split("-")
    newName = 'new-' + partName[-1]
    doc.save(newName)
    print("Added successfully!")

else:
    print("catalog file must be .txt !")
    exit()