Add support for customizable file extensions in config and file processing

config.example.yml

@@ -0,0 +1,14 @@
+# API KEY
+api_key: sk-xxxxxxx
+base_url: https://api.openai.com/v1
+# 设置默认模型。需采用OpenAI格式，其他模型可能需要使用 One-API 或 New-API 转换。
+model: gpt-4o
+# 请求时使用的代理地址
+# proxy: http://localhost:7890
+# TMDB API Key，可以在 https://www.themoviedb.org/settings/api 获取，设置后只要给定 tmdb-id 就能从 TMDB 获取数据然后传给GPT。
+# tmdb_api_key: eyxxxxxx
+# TMDB 返回结果语言
+tmdb_language: zh-CN
+# 要重命名的文件的扩展名
+exts: .mkv;.mp4;.srt;.ass
+#exts: [.mkv, .mp4, .srt]

gpt_shows_rename/__main__.py

@@ -7,7 +7,7 @@ import asyncio
 
 async def main():
     cfg = load_config()
-    files = gen_input_list(cfg.input)
+    files = gen_input_list(cfg.input, cfg.exts)
     tmdb_data = None
     if cfg.tmdb_id and not cfg.no_tmdb:
         if not cfg.tmdb_api_key:

gpt_shows_rename/config.py

@@ -1,9 +1,13 @@
 import argparse
 import openai
 import os.path
-from typing import Optional
+from typing import List, Optional
 import yaml
 
+
+DEFAULT_EXTS = ['.mp4', '.mkv', '.ass', '.srt']
+
+
 class Config:
     def __init__(self, args: argparse.Namespace, yaml_config: dict):
         self._args = args
@@ -20,6 +24,21 @@ class Config:
     def base_url(self) -> Optional[str]:
         return self._args.base_url if self._args.base_url is not None else self._yaml_config.get('base_url', 'https://api.openai.com/v1')
 
+    @property
+    def exts(self) -> List[str]:
+        if self._args.exts:
+            return self._args.exts
+        exts = self._yaml_config.get('exts')
+        if isinstance(exts, str):
+            return exts.split(";")
+        if isinstance(exts, list):
+            aexts = []
+            for ext in exts:
+                if isinstance(ext, str):
+                    aexts.append(ext)
+            return aexts
+        return DEFAULT_EXTS
+
     @property
     def hardlink(self) -> bool:
         return self._args.hardlink
@@ -91,6 +110,7 @@ def get_arg_parser() -> argparse.ArgumentParser:
     parser.add_argument('-H', '--hardlink', action='store_true', help='Use hardlink instead of symlink.')
     parser.add_argument('-n', '--no-tmdb', action='store_true', help='Do not use TMDB API to obtain data.')
     parser.add_argument('-S', '--season-number', type=int, help='Season number (optional)')
+    parser.add_argument('-e', '--exts', action='append', help='File extensions to process (optional)')
     parser.add_argument('input', help='Input directory.')
     parser.add_argument('output', help='Output directory.')
     return parser

gpt_shows_rename/file.py

@@ -4,10 +4,7 @@ from typing import List
 from .gpt import Files
 
 
-EXTS = ['.mp4', '.mkv', '.ass', '.srt']
+def gen_input_list(dir: str, exts: List[str], prefix: str = None) -> List[str]:
-
-
-def gen_input_list(dir: str, prefix: str = None) -> List[str]:
     if prefix is None:
         prefix = dir
     re = []
@@ -16,11 +13,11 @@ def gen_input_list(dir: str, prefix: str = None) -> List[str]:
             continue
         path = os.path.join(dir, f)
         if os.path.isdir(path):
-            data = gen_input_list(path, prefix=prefix)
+            data = gen_input_list(path, exts, prefix)
             re += data
         else:
-            exts = os.path.splitext(f)[1]
+            ext = os.path.splitext(f)[1]
-            if exts not in EXTS:
+            if ext not in exts:
                 continue
             p = os.path.relpath(path, prefix)
             re.append(p)