目录下查找导入特定库的文件

目录下查找导入特定库的文件

import os


def find_import_module(file_dir_path, txt_path, module_name='site'):
  with open(txt_path, 'w', encoding='utf-8') as f:
    # 遍历目录及其子目录
    for root, _, files in os.walk(file_dir_path):
        for file in files:
          if file.endswith('.py'):
            file_path = os.path.join(root, file)
            with open(file_path, 'r', encoding='utf-8') as ff:
              file_name = os.path.basename(file_path)
              for line in ff.readlines():
                # if f"import {module_name}" in line and line.strip()[0] != '#':
                if f"{module_name}" in line and line.strip()[0] != '#':
                  # print(f"Found import {module_name} in {file_path}")
                  f.write("".join((file_name, "\n")))
                  break
                if 'def' in line:
                  break

if __name__ == '__main__':
  module_name = 'torch'
  txt_path = f"./{module_name}.txt"
  file_dir_path = "./test_dir"
  find_import_module(file_dir_path, txt_path, module_name=module_name)