Эх сурвалжийг харах

refactor: enhance argument parsing in main function to support additional keyword arguments

myhloli 4 сар өмнө
parent
commit
b943f04e15
1 өөрчлөгдсөн 40 нэмэгдсэн , 1 устгасан
  1. 40 1
      mineru/cli/client.py

+ 40 - 1
mineru/cli/client.py

@@ -10,6 +10,7 @@ from ..version import __version__
 from .common import do_parse, read_fn, pdf_suffixes, image_suffixes
 
 @click.command(context_settings=dict(ignore_unknown_options=True, allow_extra_args=True))
+@click.pass_context
 @click.version_option(__version__,
                       '--version',
                       '-v',
@@ -138,11 +139,49 @@ from .common import do_parse, read_fn, pdf_suffixes, image_suffixes
 
 
 def main(
+        ctx,
         input_path, output_dir, method, backend, lang, server_url,
         start_page_id, end_page_id, formula_enable, table_enable,
         device_mode, virtual_vram, model_source, **kwargs
 ):
 
+    # 解析额外参数
+    extra_kwargs = {}
+    i = 0
+    while i < len(ctx.args):
+        arg = ctx.args[i]
+        if arg.startswith('--'):
+            param_name = arg[2:].replace('-', '_')  # 转换参数名格式
+            i += 1
+            if i < len(ctx.args) and not ctx.args[i].startswith('--'):
+                # 参数有值
+                try:
+                    # 尝试转换为适当的类型
+                    if ctx.args[i].lower() == 'true':
+                        extra_kwargs[param_name] = True
+                    elif ctx.args[i].lower() == 'false':
+                        extra_kwargs[param_name] = False
+                    elif '.' in ctx.args[i]:
+                        try:
+                            extra_kwargs[param_name] = float(ctx.args[i])
+                        except ValueError:
+                            extra_kwargs[param_name] = ctx.args[i]
+                    else:
+                        try:
+                            extra_kwargs[param_name] = int(ctx.args[i])
+                        except ValueError:
+                            extra_kwargs[param_name] = ctx.args[i]
+                except:
+                    extra_kwargs[param_name] = ctx.args[i]
+            else:
+                # 布尔型标志参数
+                extra_kwargs[param_name] = True
+                i -= 1
+        i += 1
+
+    # 将解析出的参数合并到kwargs
+    kwargs.update(extra_kwargs)
+
     if not backend.endswith('-client'):
         def get_device_mode() -> str:
             if device_mode is not None:
@@ -188,7 +227,7 @@ def main(
                 table_enable=table_enable,
                 server_url=server_url,
                 start_page_id=start_page_id,
-                end_page_id=end_page_id
+                end_page_id=end_page_id,
                 **kwargs,
             )
         except Exception as e: