From a463548731a1fcc95706192f3872dc38f4aac339 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=B2=9B=E9=A3=8E?= <i@loli.xin>
Date: Thu, 27 Jun 2024 01:17:57 +0800
Subject: [PATCH] =?UTF-8?q?=E6=B7=BB=E5=8A=A0=E5=AE=9A=E5=88=B6=E8=BE=93?=
 =?UTF-8?q?=E5=87=BA=E5=8A=9F=E8=83=BD?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 urls/main.py | 19 +++++++++++--------
 1 file changed, 11 insertions(+), 8 deletions(-)

diff --git a/urls/main.py b/urls/main.py
index 9dfaea1..2b1c9aa 100644
--- a/urls/main.py
+++ b/urls/main.py
@@ -3,6 +3,7 @@ import urllib.parse
 import argparse
 import re
 import time
+import yaml
 
 def natural_sort_key(s):
     """为了自然排序的辅助函数，将字符串中的数字部分转换成整数"""
@@ -29,25 +30,27 @@ def generate_urls(file_paths, base_url, sub_directory, min_size):
         urls[dir_name].append(f"{file_size_bytes}:{current_timestamp}:{url}")
     return urls
 
-def save_urls(urls, output_file):
-    """将URL链接保存到文本文件中"""
+def save_urls(urls, output_file, root_folder):
+    """将URL链接保存到YAML文件中"""
     with open(output_file, 'w', encoding='utf-8') as f:
+        f.write(f"{root_folder}:\n")
         for dir_name, files in urls.items():
             if dir_name == '.':
-                for file in files:
-                    f.write(f"{file}\n")
-            else:
-                f.write(f"{dir_name}:\n")
                 for file in files:
                     f.write(f"  {file}\n")
+            else:
+                f.write(f"  {dir_name}:\n")
+                for file in files:
+                    f.write(f"    {file}\n")
 
 def parse_arguments():
     """解析命令行参数"""
     parser = argparse.ArgumentParser(description='Generate URLs from file names.')
     parser.add_argument('--dir', type=str, default='', help='Sub-directory for generating file URLs (optional)')
-    parser.add_argument('--output', type=str, default='urls.txt', help='Output file name (default: urls.txt)')
+    parser.add_argument('--output', type=str, default='urls.yaml', help='Output file name (default: urls.yaml)')
     parser.add_argument('--base-url', type=str, default='https://link.kite.kim/feng', help='Base URL for generating file URLs (default: https://link.kite.kim/feng)')
     parser.add_argument('--min-size', type=int, default=128*1024, help='Minimum file size in bytes (default: 128KB)')
+    parser.add_argument('--rf', type=str, required=True, help='Root folder name for YAML output')
     return parser.parse_args()
 
 def list_files_recursive(start_path='.', exclude_files=None):
@@ -67,7 +70,7 @@ def main():
     exclude_files = {current_script}  # 排除当前脚本文件
     file_paths = list_files_recursive('.', exclude_files)
     urls = generate_urls(file_paths, args.base_url, args.dir, args.min_size)
-    save_urls(urls, args.output)
+    save_urls(urls, args.output, args.rf)
     print(f"URL链接已保存到{args.output}")
 
 if __name__ == '__main__':