Benchmark code formatting, more refactoring

Author: scotts

benchmarks/decoders/benchmark_decoders.py

@@ -31,7 +31,7 @@
 class DecoderKind:
     display_name: str
     kind: typing.Type[AbstractDecoder]
-    default_options: dict = field(default_factory=dict)
+    default_options: dict[str, str] = field(default_factory=dict)
 
 
 decoder_registry = {
@@ -66,6 +66,16 @@ def get_test_resource_path(filename: str) -> str:
     return str(Path(__file__).parent / f"../../test/resources/{filename}")
 
 
+def parse_options_code(options_code: str) -> dict[str, str]:
+    options = {}
+    for item in options_code.split("+"):
+        if item.strip() == "":
+            continue
+        k, v = item.split("=")
+        options[k] = v
+    return options
+
+
 def main() -> None:
     """Benchmarks the performance of a few video decoders"""
 
@@ -98,19 +108,17 @@ def main() -> None:
         "--decoders",
         help=(
             "Comma-separated list of decoders to benchmark. "
-            "Choices are: "
-            + ", ".join(decoder_registry.keys())
-            + ". "
-            + "To specify options, append a ':' and then value pairs seperated by a '+'. "
-            "For example, torchcodec:num_threads=1+color_conversion_library=filtergraph."
+            "Choices are: " + ", ".join(decoder_registry.keys()) + ". "
+            "To specify options, append a ':' and then value pairs seperated by a '+'. "
+            "For example, torchcodec_core:num_threads=1+color_conversion_library=filtergraph."
         ),
         type=str,
         default=(
             "decord,decord_batch,"
-            + "torchvision,"
-            + "torchaudio,"
-            + "torchcodec_core,torchcodec_core:num_threads=1,torchcodec_core_batch,torchcodec_core_nonbatch,"
-            + "torchcodec_public"
+            "torchvision,"
+            "torchaudio,"
+            "torchcodec_core,torchcodec_core:num_threads=1,torchcodec_core_batch,torchcodec_core_nonbatch,"
+            "torchcodec_public"
         ),
     )
     parser.add_argument(
@@ -135,27 +143,19 @@ def main() -> None:
     decoders_to_run = {}
     for decoder in specified_decoders:
         if ":" in decoder:
-            decoder_name, _, options = decoder.partition(":")
-            assert decoder_name in decoder_registry
-
-            kwargs_dict = {}
-            for item in options.split("+"):
-                if item.strip() == "":
-                    continue
-                k, v = item.split("=")
-                kwargs_dict[k] = v
-
-            display_name = decoder_registry[decoder_name].display_name
-            kind = decoder_registry[decoder_name].kind
-            decoders_to_run[display_name + options] = kind(**kwargs_dict)
+            decoder, _, options_code = decoder.partition(":")
+            assert decoder in decoder_registry
+            display = decoder_registry[decoder].display_name + options_code
+            options = parse_options_code(options_code)
         elif decoder in decoder_registry:
-            display_name = decoder_registry[decoder].display_name
-            kind = decoder_registry[decoder].kind
-            default_options = decoder_registry[decoder].default_options
-            decoders_to_run[display_name] = kind(**default_options)
+            display = decoder_registry[decoder].display_name
+            options = decoder_registry[decoder].default_options
         else:
             raise ValueError(f"Unknown decoder: {decoder}")
 
+        kind = decoder_registry[decoder].kind
+        decoders_to_run[display] = kind(**options)
+
     video_paths = args.bm_video_paths.split(",")
     if args.bm_video_dir:
         video_paths = []