在指定目录下的文件中搜索字符串

在指定目录下的文件中搜索字符串

程序参数为： search [dir] [regex]
例如：

search document “^#” ： 在 ./document 目录下的所有文件中查找符合正则表达式 ^# 的文件

代码比较简单，使用了线程池以加速搜索

import logging
import os
import re
import sys
from concurrent.futures.thread import ThreadPoolExecutor

import chardet

# main.py dir str
# regex support

# 创建线程池
exector = ThreadPoolExecutor()


def is_binary_file(file_path: str):
    binary_array = (
        'png', 'jpg', 'ico', 'icns'                         # 图片
        , 'woff2', 'ttf', 'otf'                             # 字体
        , 'zip', '7z', 'rar', 'gz', 'xz', 'tar'             # 压缩文件
        , 'mp3', 'wav', 'ape', 'flac'                       # 音乐文件
        , 'mp4', 'avi', 'flv'                               # 音频文件
        , 'docx', 'doc', 'xlsx', 'xls', 'pdf'               # 文档
        , 'obj', 'exe', 'idx', 'bin', 'lib', 'a', 'dll'     # 其他二进制文件
        , 'dat', 'iso', 'assets', 'resource', 'pack', 'pak'
    )
    logging.debug(file_path.split('.')[-1])
    if file_path.split('.')[-1] in binary_array:
        return True
    return False


def search_in_file(file_path: str):
    # 检测文件编码
    data: bytes = bytes(10)
    with open(file_path, 'rb') as f:
        data = f.read()
    encoding_result = chardet.detect(data)
    # 搜索文件
    try:
        file_content: str = data.decode(encoding_result['encoding'])
        result = re.search(sys.argv[2], file_content)
        if result is not None:
            print(file_path)
    except Exception:
        logging.error("%s 查看失败", file_path)


def search_from_dir() -> None:
    dir_name = os.path.join(os.getcwd(), sys.argv[1])
    for dirpath, dirnames, filenames in os.walk(dir_name):
        for file_name in filenames:
            file_path = os.path.join(dirpath, file_name)

            # 判断是否是二进制文件
            if is_binary_file(file_path):
                continue
            # 若文件太大就跳过
            file_size = int(os.stat(file_path).st_size/1024/1024)  # 单位为 MB
            if file_size > 8:
                logging.warning("文件 %s 由于尺寸过大而被跳过", file_path)
                continue
            # 将任务提交给线程池
            exector.submit(search_in_file, file_path)


if __name__ == "__main__":
    logging.basicConfig(level=logging.WARNING)
    if len(sys.argv) != 3:
        sys.exit(-1)
    search_from_dir()
    # 关闭线程池
    exector.shutdown(wait=True)