Cách tìm kiếm chuỗi trong nhiều định dạng tài liệu (CSV, Text, MS Word) với Python

Vấn đề

Giả sử bạn có một thư mục chứa nhiều file với các định dạng khác nhau, cần tìm một từ khóa cụ thể trong tất cả các file đó.

Chuẩn bị

Cài đặt các gói sau:

1. beautifulsoup4
2. python-docx

Cách thực hiện

1. Hàm tìm chuỗi trong file CSV

Sử dụng module csv.reader để duyệt qua file và tìm chuỗi, trả về True nếu tìm thấy, ngược lại False.

def csv_stringsearch(input_file, input_string):
    """ Tìm chuỗi trong file CSV.
    args: input_file, input_string """
    with open(input_file) as file:
        for row in csv.reader(file):
            for column in row:
                if input_string in column.lower():
                    return True
    return False

2. Hàm tìm chuỗi trong file văn bản (Text)

Đọc file văn bản phức tạp hơn do phải xử lý encoding. Có hàng nghìn loại encoding và xác định encoding chính là phần khó khăn nhất. Tự động hóa bằng cách sử dụng UnicodeDammit để đoán encoding.

def text_stringsearch(input_file, input_string):
    """ Tìm chuỗi trong file text.
    args: input_file, input_string """
    with open(input_file, 'rb') as file:
        content = file.read(1024)
        guessencoding = UnicodeDammit(content)
        encoding = guessencoding.original_encoding
    
    # Mở và đọc với encoding đã xác định
    with open(input_file, encoding=encoding) as file:
        for line in file:
            if input_string in line.lower():
                return True
    return False

3. Hàm tìm chuỗi trong tài liệu MS Word (.docx)

def MSDocx_stringsearch(input_file, input_string):
    """ Tìm chuỗi trong tài liệu MS Word.
    args: input_file, input_string """
    doc = docx.Document(input_file)
    for paragraph in doc.paragraphs:
        if input_string in paragraph.text.lower():
            return True
    return False

4. Hàm chính (main) để duyệt qua các file

Lặp qua các file trong thư mục và gọi hàm tương ứng với đuôi file. Giả sử code và file đầu vào nằm cùng thư mục.

def main(input_string):
    """ Mở thư mục hiện tại và tìm chuỗi trong tất cả file.
    args: input_string """
    for root, dirs, files in os.walk('.'):
        for file in files:
            # Lấy đuôi file
            extension = file.split('.')[-1]
            if extension in function_mapping:
                search_file = function_mapping.get(extension)
                full_file_path = os.path.join(root, file)
                if search_file(full_file_path, input_string):
                    print(f' *** Tìm thấy chuỗi trong {full_file_path}')

5. Ánh xạ hàm với đuôi file bằng dictionary

function_mapping = {
    'csv': csv_stringsearch,
    'txt': text_stringsearch,
    'docx': MSDocx_stringsearch,
}

6. Chương trình hoàn chỉnh

import os
import argparse
import csv
import docx
from bs4 import UnicodeDammit

def csv_stringsearch(input_file, input_string):
    """ Tìm chuỗi trong file CSV. """
    with open(input_file) as file:
        for row in csv.reader(file):
            for column in row:
                if input_string in column.lower():
                    return True
    return False

def MSDocx_stringsearch(input_file, input_string):
    """ Tìm chuỗi trong tài liệu MS Word. """
    doc = docx.Document(input_file)
    for paragraph in doc.paragraphs:
        if input_string in paragraph.text.lower():
            return True
    return False

def text_stringsearch(input_file, input_string):
    """ Tìm chuỗi trong file text. """
    with open(input_file, 'rb') as file:
        content = file.read(1024)
        guessencoding = UnicodeDammit(content)
        encoding = guessencoding.original_encoding
    with open(input_file, encoding=encoding) as file:
        for line in file:
            if input_string in line.lower():
                return True
    return False

def main(input_string):
    """ Duyệt thư mục và tìm chuỗi trong tất cả file. """
    for root, dirs, files in os.walk('.'):
        for file in files:
            extension = file.split('.')[-1]
            if extension in function_mapping:
                search_file = function_mapping.get(extension)
                full_file_path = os.path.join(root, file)
                if search_file(full_file_path, input_string):
                    print(f' *** Tìm thấy chuỗi trong {full_file_path}')

function_mapping = {
    'csv': csv_stringsearch,
    'txt': text_stringsearch,
    'docx': MSDocx_stringsearch,
}

if __name__ == '__main__':
    string_to_search = 'Hello'
    print(f'Output \n')
    main(string_to_search.lower())

Kết quả mẫu:

*** Tìm thấy chuỗi trong .\Hello_World.docx
*** Tìm thấy chuỗi trong .\My_Amazing_WordDoc.docx

7. Chạy từ dòng lệnh với argparse

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('-s', type=str, help='Chuỗi cần tìm', default='Hello')
    args = parser.parse_args()
    main(args.s.lower())