11
\$\begingroup\$

I have the following Python script. The purpose of the script is to extract text from PDFs. I use textract for that because soon I realized there is no easy way to check if a page contains an image or not. So I extract the whole text using textract.

The workflow is like this. main() parses each pdf file from a folder, I extract the text, I search for keyword strikes and then I export the result to a csv file inside folder output_results.

What I can fix on my logic? What can I change in my code? I find it messy, how I can clean it up?

import textract
import os
import csv

class PdfMiner():

    path = os.getcwd() + '/folderForPdf/'
    output_path = os.getcwd() + '/output_results/'

    def __init__(self):
        pass

    def main(self):
        for self.filename in os.listdir(self.path):
            self.text = (PdfMiner().extract_text_from_pdf(self.path + self.filename))
            self.keyword_strike_dict = PdfMiner().keyword_strike(self.text)
            if bool(self.keyword_strike_dict):
                PdfMiner().output_to_csv(self.filename, self.keyword_strike_dict)


    def keyword_strike(self, text, keyword_strike_dict={}):
        '''keyword_strike function counts how many times a specific keyword occurs'''
        self.keyword_strike_dict = {}
        self.text = text
        self.keywords_list = PdfMiner().extract_keywords()
        for keyword in self.keywords_list:
            if keyword in text.decode('utf-8'):
                self.keyword_strike_dict[keyword] = text.decode('utf-8').count(keyword)
        return self.keyword_strike_dict

    def extract_keywords(self, keywords_list=None):
        '''function extract_keywords extract the keywords from file keywords.txt, into a list'''
        keywords_list = []
        with open('keywords.txt', 'r', encoding='utf8') as keywords_file:
            for keyword in keywords_file:
                keywords_list.append(keyword.strip('\n'))
        return keywords_list

    def extract_text_from_pdf(self, file_destination, text=None):
        '''extract_text_from_pdf'''
        self.file_destination = file_destination
        text = textract.process(self.file_destination, method='tesseract', language='eng', encoding='utf-8')
        return text

    def output_to_csv(self, *args, **kwargs):
        '''output_csv exports results to csv'''
        self.filename = args[0]
        self.keyword_strike_dict = args[1]
        self.output_file_path = PdfMiner().output_path + self.filename.strip('.pdf')
        with open(self.output_file_path + '.csv', 'w+', newline='') as csvfile:
            row_writer = csv.writer(csvfile, delimiter=',')
            row_writer.writerow(['keyword', 'keyword_count'])
            for keyword, keyword_count in self.keyword_strike_dict.items():
                print(keyword, keyword_count)
                row_writer.writerow([keyword, keyword_count])

if __name__ == "__main__":
    PdfMiner().main()
\$\endgroup\$
11
\$\begingroup\$

I don't see a good reason why this should be a class. You only have two things in your state, self.text, which you could pass as an argument, and self.path, self.output_path, which I would also pass as arguments, maybe with a default value.

Also, you are probably using classes wrong if your class has a main method that needs to instantiate new instances of the class on the fly.

Your algorithm is not very efficient. You need to run over the whole text twice for each keyword. Once to check if it is in there and then again to count it. The former is obviously redundant, since str.count will just return 0 if the value is not present.

However, what would be a better algorithm is to first extract all the words (for example using a regex filtering only letters) and then count the number of times each word occurs using a collections.Counter, optionally filtering it down to only those words which are keywords. It even has a most_common method, so your file will be ordered by number of occurrences, descending.

Instead of mucking around with os.getcwd() and os.listdir, I would recommend to use the (Python 3) pathlib.Path object. It supports globbing (to get all files matching a pattern), chaining them to get a new path and even replacing the extension with a different one.

When reading the keywords, you can use a simple list comprehension. Or, even better, a set comprehension to get in calls for free.

line.strip() and line.strip("\n") are probably doing the same thing, unless you really want to preserve the spaces at the end of words.

At the same time, doing self.filename.strip('.pdf') is a bit dangerous. It removes all characters given, until none of the characters is found anymore. For example, "some_file_name_fdp.pdf" will be reduced to "some_file_name_".

The csv.writer has a writerows method that takes an iterable of rows. This way you can avoid a for loop.

I would ensure to run only over PDF files, otherwise you will get some errors if a non-PDF file manages to sneak into your folder.

I have done all of this in the following code (not tested, since I don't have textract installed ATM):

from collections import Counter
import csv
from pathlib import Path
import re
import textract

def extract_text(file_name):
    return textract.process(file_name, method='tesseract', language='eng',
                            encoding='utf-8').decode('utf-8')

def extract_words(text):
    return re.findall(r'([a-zA-Z]+)', text)

def count_keywords(words, keywords):
    return Counter(word for word in words if word in keywords)

def read_keywords(file_name):
    with open(file_name) as f:
        return {line.strip() for line in f}

def save_keywords(file_name, keywords):
    with open(file_name, "w", newline='') as csvfile:
        writer = csv.writer(csvfile, delimiter=',')
        writer.writerow(['keyword', 'keyword_count'])
        writer.writerows(keywords.most_common())

def main():
    output_folder = Path("output_results")
    keywords = read_keywords('keywords.txt')

    for f in Path("folderForPdf").glob("*.pdf"):
        words = extract_words(extract_text(f))
        keyword_counts = count_keywords(words, keywords)
        save_keywords(output_folder / f.with_suffix(".csv"), keyword_counts)

if __name__ == "__main__":
    main()
\$\endgroup\$
  • \$\begingroup\$ Wow amazing job thank you very much. I see, most of my code is pointless \$\endgroup\$ – Iakovos Belonias May 20 at 10:33
  • 2
    \$\begingroup\$ @IakovosBelonias Not pointless (it worked before, didn't it?), just a bit too verbose, maybe ;) \$\endgroup\$ – Graipher May 20 at 10:35
  • 1
    \$\begingroup\$ Thank you very much \$\endgroup\$ – Iakovos Belonias May 20 at 10:36
  • \$\begingroup\$ Should I use regex considering that I don't care to much about extract the text but mostly to count the number of occurrence? \$\endgroup\$ – Iakovos Belonias May 20 at 10:48
  • \$\begingroup\$ @IakovosBelonias Well, with this approach you need to find the words of the text first in order to count them. This is one advantage of your approach, at the cost of performance and false positives in case of partial matches with the keyword. \$\endgroup\$ – Graipher May 20 at 10:52

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.