Attempting an OOP approach for PDF Paranoia - Automate the Boring Stuff CH 15

Question

This exercise is called PDF Paranoia and it comes from CH 15 of the book Automate the Boring Stuff With Python - second edition. There are two parts to the exercise, which I've separated into two Python files. The code I've provided uses Python 3.10 and PyPDF2 v. 2.10.0.

The instructions are as follows:

Using the os.walk() function from Chapter 10, write a script that will go through every PDF in a folder (and its subfolders) and encrypt the PDFs using a password provided on the command line. Save each encrypted PDF with an _encrypted.pdf suffix added to the original filename. Before deleting the original file, have the program attempt to read and decrypt the file to ensure that it was encrypted correctly.

Then, write a program that finds all encrypted PDFs in a folder (and its subfolders) and creates a decrypted copy of the PDF using a provided password. If the password is incorrect, the program should print a message to the user and continue to the next PDF.

My three main goals with the code are:

1. Use an object-oriented approach
2. Write useful docstrings / comments
3. Write Pythonic code

I also implemented storing functions as a callable (without the "()") in order to reuse the code. If there is a better approach for this or if there is something I have missed, I would love to know.

First file, stored as pdfparanoia.py:

#!python3
# pdfparanoia.py - Walk through a directory and its subdirectories, and encrypt
#                  every PDF file. Provide the encryption password via command
#                  line and the filepath to search through. Check to see if
#                  encryption succeeded by attempting to read and decrypt the
#                  files. Prompt the user to delete the older unencrypted files.
#
#           USAGE: sys.argv[1] = password
#                  sys.argv[2] = filepath to search for PDFs
#                  Using Python 3.10 and PyPDF2 v2.10.0


import os
import sys
import PyPDF2
from typing import Callable


class EncryptPDFs:
    def __init__(self):
        self.password = sys.argv[1]
        self.dir_to_search = sys.argv[2]
        self.encryption_call = self.encrypt_pdf # store as callable
        self.delete_call = self.delete_file # store as callable
        self.encryption_check_call = self.encryption_checker # store as callable

    @staticmethod
    def search_for_pdf(dir_to_search: str, function_call: Callable) -> None:
        """
        :param dir_to_search: input directory tree to walk through
        :param function_call: Input callable as argument and pass it the file variable.
        Callable argument acts on all PDF files in the given directory tree.
        """
        print(f"Searching for .pdf files in: {dir_to_search}")
        for folder, sub_folders, file_list in os.walk(dir_to_search):
            for file in file_list:
                if file.split(".")[-1].lower() == "pdf":
                    function_call(file)

    def encrypt_pdf(self, pdf_file: str) -> None:
        """
        :param file: Input PDF file name from the calling function (search_for_pdf()).
        If the PDF is not encrypted, create a PDFFileWriter object to encrypt
        and save it as a new file. Encrypted files will be skipped.
        """
        file_object = open(pdf_file, "rb")
        reader = PyPDF2.PdfFileReader(file_object)
        if not reader.isEncrypted:
            print(f"Encrypting {pdf_file}")
            writer = PyPDF2.PdfFileWriter()
            for page_num in range(reader.numPages):
                writer.addPage(reader.getPage(page_num))
            writer.encrypt(self.password)
            result_pdf = open(pdf_file.split(".")[0] + "_encrypted.pdf", "wb")
            writer.write(result_pdf)
            result_pdf.close()
        file_object.close()


    def encryption_checker(self, pdf_file: str) -> None:
        """Check to see if encryption succeeded by reading and decrypting the files."""
        file_object = open(pdf_file, "rb")
        reader = PyPDF2.PdfFileReader(file_object)
        try:
            reader.getPage(0)
        except PyPDF2.errors.FileNotDecryptedError:
            print(f"Encryption check 1: {pdf_file} is encrypted and cannot be read.")
        if reader.isEncrypted:
            print(
                f"Encryption check 2: isEncrypted() method returns True when called on {pdf_file}."
            )
            print(f"{pdf_file} is encrypted.\n")
            reader.decrypt(self.password)
        file_object.close()

    def delete_file(self, pdf_file: str) -> None:
        """Optionally delete unencrypted files after encryption check."""
        file_object = open(pdf_file, "rb")
        reader = PyPDF2.PdfFileReader(file_object)
        if not reader.isEncrypted:
            file_object.close()
            os.remove(pdf_file)

    def main(self) -> None:
        """
        Search for PDF files, prompt user to enter encryption password, then encrypt
        the files. Attempt to read / decrypt the files to ensure encryption succeeded,
        then prompt user to delete unencrypted files and exit.
        """
        self.search_for_pdf(self.dir_to_search, self.encryption_call)
        print("\nSearching for encrypted files to check.\n")
        self.search_for_pdf(self.dir_to_search, self.encryption_check_call)
        while True:
            choice = input("\nDo you wish to delete all unencrypted files? Y/N\n")
            if choice.upper() == "Y":
                self.search_for_pdf(self.dir_to_search, self.delete_call)
                sys.exit()
            elif choice.upper() == "N":
                sys.exit()
            else:
                continue


if __name__ == "__main__":
    EncryptPDFs().main()

Second file, stored as decrypt_pdf_file.py

This file imports the first file in order to call the search_for_pdf() function:

#!python3
# decrypt_pdf_file.py - Walk through a directory and its subdirectories, and prompt
#                       the user to input a password to decrypt the files.
#
#                USAGE: sys.argv[1] = filepath to search for PDFs


import sys
import PyPDF2
from pdfparanoia import EncryptPDFs

dir_to_search = sys.argv[1]


def decrypt_pdf_file(file: str) -> None:
    """
    :param file: input PDF file from calling function. Create a PdfFileReader
    object to determine if the PDF is encrypted. If the PDF is encrypted, prompt
    the user for the encryption password. If the password is correct, decrypt it
    the file and save it. If the password is incorrect, warn the user and continue.
    """
    file_object = open(file, 'rb')
    reader = PyPDF2.PdfFileReader(file_object)
    if reader.isEncrypted:
        password = input(f'Please enter the password to decrypt {file}\n')
        if reader.decrypt(password) == 0:
            print(f'Password invalid. {file} was not decrypted\n')
        elif reader.decrypt(password) == 1:
            print(f'Decrypting {file}')
            writer = PyPDF2.PdfFileWriter()
            for page_num in range(reader.numPages):
                writer.addPage(reader.getPage(page_num))
            result_pdf = open(file.split('.')[0] + '_decrypted.pdf', 'wb')
            writer.write(result_pdf)
            result_pdf.close()
            print(f'{file} has been decrypted and saved.')
    file_object.close()


DECRYPTION_CALL = decrypt_pdf_file  # store function as Callable

if __name__ == "__main__":
    EncryptPDFs.search_for_pdf(dir_to_search, DECRYPTION_CALL)

Answer 1

Object oriented code means that you have objects, and methods for interacting with those objects, you can then very easily reuse / compose those objects to do different things.

In this case you've created a god class (https://en.wikipedia.org/wiki/God_object) which does several different things (reading arguments from the command line; finding PDFs; encrypting / decrypting PDFs; deleting PDFs).

Ideally each of those concerns should be a separate entity in code (https://en.wikipedia.org/wiki/Separation_of_concerns). In traditional OO languages that would mean either separate strategy classes (normally used when you have multiple related behaviours which are tightly coupled, like encrypting and decrypting) or through using 'template methods' for those parts which can be overridden through inheritance (used to change one step).

In python you generally don't use template methods, for many reasons, instead you pass in function objects to override specific steps in an algorithm. You've kind of done that here, but you've gotten confused over how variables work in python:

    self.delete_call = self.delete_file # store as callable

^ This code doesn't store anything, or achieve anything really. Variables in python are just labels, you've just added an alias to 'self.delete_file'. It would be better to just pass in 'self.delete_file' directly where it's used.

Also, OO code still normally has a high level main function which kicks everything off. In pure OO languages like Java this is a static method of a class, but that's not needed in python, you can just write a function, writing staticmethods is very much 'not pythonic'.

When you have to do something, and then clean it up at the end (i.e. opening and closing files in your code) - the pythonic way of doing that is to use a context manager. That is, instead of:

file_object = open(file, 'rb')
# Code here
file_object.close()

You would do:

with open(file, 'rb') as file_object:
    # Code here

This is neater, but also more robust, the closing of the file will happen now even if your code crashes. Although, in this case it's also not completely necessary, the library you're using supports creating readers / writers from strings directly, and supports using them as context managers. In addition, it supports iterating over pages directly, which will also be more efficient and more pythonic:

def encrypt_pdf(self, pdf_file: str) -> None:
    """
    :param file: Input PDF file name from the calling function (search_for_pdf()).
    If the PDF is not encrypted, create a PDFFileWriter object to encrypt
    and save it as a new file. Encrypted files will be skipped.
    """
    with PyPDF2.PdfFileReader(pdf_file) as reader:
        if not reader.isEncrypted:
            print(f"Encrypting {pdf_file}")
            with PyPDF2.PdfFileWriter(pdf_file.split('.')[0] + '_encrypted.pdf') as writer:
                for page in reader.pages:
                    writer.addPage(page)
                writer.encrypt(self.password)

Some of your methods do too much. Your method for finding PDFs should really just find PDFs, and leave doing stuff to them for something else. I.e, not this:

@staticmethod
def search_for_pdf(dir_to_search: str, function_call: Callable) -> None:
    """
    :param dir_to_search: input directory tree to walk through
    :param function_call: Input callable as argument and pass it the file variable.
    Callable argument acts on all PDF files in the given directory tree.
    """
    print(f"Searching for .pdf files in: {dir_to_search}")
    for folder, sub_folders, file_list in os.walk(dir_to_search):
        for file in file_list:
            if file.split(".")[-1].lower() == "pdf":
                function_call(file)

...

self.search_for_pdf(self.dir_to_search, self.encryption_call)

Rather, this:

def search_for_pdf(dir_to_search: str) -> None:
    """
    :param dir_to_search: input directory tree to walk through
    :param function_call: Input callable as argument and pass it the file variable.
    Callable argument acts on all PDF files in the given directory tree.
    """
    print(f"Searching for .pdf files in: {dir_to_search}")
    for folder, sub_folders, file_list in os.walk(dir_to_search):
        for file in file_list:
            if file.split(".")[-1].lower() == "pdf":
                yield file

...

for pdf in search_for_pdf(self.dir_to_search):
    encrypt(pdf)

As there isn't much logic here, you could quite comfortably handle the whole thing without any classes at all, and indeed that would probably be better in this case. Object-oriented is useful when you have multiple strategies that you want to apply and the strategies are complex enough to need multiple methods. Here you can just really just use functions. That said, to give you what you want, let's try splitting our four concerns into 3 functions and a class:

import os
import sys
import PyPDF2


class PdfEncryptor:
    def __init__(self, password: str):
        self._password = password
        
    def encrypt(self, pdf_writer):
        pdf_writer.encrypt(self._password)
    
    def decrypt(self, pdf_reader):
        pdf_reader.decrypt(self._password)
    
    def encrypt_path(self, pdf_path):
        with PyPDF2.PdfFileReader(pdf_path) as reader:
            if not reader.isEncrypted:
                print(f"Encrypting {pdf_path}")
                with PyPDF2.PdfFileWriter(pdf_path.split('.')[0] + '_encrypted.pdf') as writer:
                    for page in reader.pages:
                        writer.addPage(page)
                    self.encrypt(writer)
    
    def check_can_decrypt(self, pdf_path):
        """Check to see if encryption succeeded by reading and decrypting the files."""
        with PyPDF2.PdfFileReader(pdf_path) as reader:
            try:
                reader.getPage(0)
            except PyPDF2.errors.FileNotDecryptedError:
                print(f"Encryption check 1: {pdf_file} is encrypted and cannot be read.")
            if reader.isEncrypted:
                print(
                    f"Encryption check 2: isEncrypted() method returns True when called on {pdf_file}."
                )
                print(f"{pdf_file} is encrypted.\n")
                self.decrypt(reader)


def search_for_pdf(dir_to_search: str) -> None:
    """
    :param dir_to_search: input directory tree to walk through
    :param function_call: Input callable as argument and pass it the file variable.
    Callable argument acts on all PDF files in the given directory tree.
    """
    print(f"Searching for .pdf files in: {dir_to_search}")
    for folder, sub_folders, file_list in os.walk(dir_to_search):
        for file in file_list:
            if file.split(".")[-1].lower() == "pdf":
                yield file


def delete_file(pdf_file: str):
    """Optionally delete unencrypted files after encryption check."""
    with PyPDF2.PdfFileReader(pdf_file) as reader:
        encryped = reader.isEncrypted
    if not encrypted:
        os.remove(pdf_file)

        
def main():
    """
    Search for PDF files, prompt user to enter encryption password, then encrypt
    the files. Attempt to read / decrypt the files to ensure encryption succeeded,
    then prompt user to delete unencrypted files and exit.
    """
    password = sys.argv[1]
    encryptor = PdfEncryptor(password)
    dir_to_search = sys.argv[2]
    
    for pdf in search_for_pdf(dir_to_search):
        encryptor.encrypt_path(pdf)
    
    print("\nSearching for encrypted files to check.\n")

    for pdf in search_for_pdf(dir_to_search):
        encryptor.check_can_decrypt(pdf)
    
    while True:
        choice = input("\nDo you wish to delete all unencrypted files? Y/N\n")
        if choice.upper() == "Y":
            for pdf in search_for_pdf(dir_to_search):
                delete_file(pdf)
            sys.exit()
        elif choice.upper() == "N":
            sys.exit()
        else:
            continue

And now we can reuse the PDF finding logic (which is the only thing being reused) in your second script:

import sys
import PyPDF2
from pdfparanoia import search_for_pdf


def decrypt_pdf_file(file: str) -> None:
    """
    :param file: input PDF file from calling function. Create a PdfFileReader
    object to determine if the PDF is encrypted. If the PDF is encrypted, prompt
    the user for the encryption password. If the password is correct, decrypt it
    the file and save it. If the password is incorrect, warn the user and continue.
    """
    with PyPDF2.PdfFileReader(file) as reader:
        if reader.isEncrypted:
            password = input(f'Please enter the password to decrypt {file}\n')
            if reader.decrypt(password) == 0:
                print(f'Password invalid. {file} was not decrypted\n')
            elif reader.decrypt(password) == 1:
                print(f'Decrypting {file}')
                with PyPDF2.PdfFileWriter(file.split('.')[0] + '_decrypted.pdf') as writer:
                    writer = PyPDF2.PdfFileWriter()
                    for page in reader.pages:
                        writer.addPage(page)
                print(f'{file} has been decrypted and saved.')


if __name__ == "__main__":
    dir_to_search = sys.argv[1]
    for pdf in search_for_pdf(dir_to_search):
        decrypt_pdf_file(pdf)

Answer 2

password provided on the command line is not something you should do even if you're told so by a specification. If someone were to use this with a valuable password, that password would be exposed on the screen and saved to the shell's history. So don't do that. Acquire the password using getpass instead.