modify scripts

aabook/tools/tools_dir.py

@@ -0,0 +1,66 @@
+import os
+import shutil
+import argparse
+import re
+
+def flatten_directory(source_dir):
+    """
+    将指定目录下的所有txt文件移动到当前目录并重命名
+
+    Args:
+        source_dir: 源目录
+    """
+
+    for root, dirs, files in os.walk(source_dir):
+        for file in files:
+            if file.endswith('.txt'):
+                src_file = os.path.join(root, file)
+                dst_file = os.path.join(source_dir, f"[{os.path.basename(root)}]_{file}")
+                print(f'move {src_file} {dst_file}')
+                shutil.move(src_file, dst_file)
+
+def unflatten_directory(source_dir):
+    """
+    将当前目录下的txt文件按照文件名中的目录信息进行分类
+
+    Args:
+        source_dir: 源目录
+    """
+
+    for file in os.listdir(source_dir):
+        if file.endswith('.txt'):
+            #dir_name, filename = file.split(']', 1)[0][1:], file.split(']', 1)[1][1:]
+            # 方法二：使用正则表达式
+            match = re.match(r"\[(.*)]_(.*)", file)
+            if match:
+                dir_name, filename = match.groups()
+
+            dst_dir = os.path.join(source_dir, dir_name)
+            dst_file = os.path.join(dst_dir, filename)
+            src_file = os.path.join(source_dir, file)
+
+            # 创建目标目录
+            os.makedirs(dst_dir, exist_ok=True)
+
+            # 移动文件
+            print(f'move {src_file} {dst_file}')
+            shutil.move(src_file, dst_file)
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='Flatten or unflatten a directory of txt files')
+    parser.add_argument('-f', '--flatten', action='store_true', help='Flatten the directory')
+    parser.add_argument('-u', '--unflatten', action='store_true', help='Unflatten the directory')
+    parser.add_argument('directory', help='The directory to process')
+    args = parser.parse_args()
+
+    if args.flatten and args.unflatten:
+        print("Please choose either --flatten or --unflatten, not both.")
+        exit(1)
+
+    if args.flatten:
+        flatten_directory(args.directory)
+    elif args.unflatten:
+        unflatten_directory(args.directory)
+    else:
+        print("Please specify either --flatten or --unflatten.")
+        exit(1)