diff --git a/nerfstudio/configs/dataparser_configs.py b/nerfstudio/configs/dataparser_configs.py index 382c42ad1..80b915151 100644 --- a/nerfstudio/configs/dataparser_configs.py +++ b/nerfstudio/configs/dataparser_configs.py @@ -16,11 +16,14 @@ Aggregate all the dataparser configs in one location. """ +from typing import TYPE_CHECKING + import tyro from nerfstudio.data.dataparsers.arkitscenes_dataparser import ( ARKitScenesDataParserConfig, ) +from nerfstudio.data.dataparsers.base_dataparser import DataParserConfig from nerfstudio.data.dataparsers.blender_dataparser import BlenderDataParserConfig from nerfstudio.data.dataparsers.dnerf_dataparser import DNeRFDataParserConfig from nerfstudio.data.dataparsers.dycheck_dataparser import DycheckDataParserConfig @@ -57,11 +60,19 @@ dataparsers = { external_dataparsers = discover_dataparsers() all_dataparsers = {**dataparsers, **external_dataparsers} -AnnotatedDataParserUnion = tyro.conf.OmitSubcommandPrefixes[ # Omit prefixes of flags in subcommands. - tyro.extras.subcommand_type_from_defaults( + +if TYPE_CHECKING: + # For static analysis (tab completion, type checking, etc), just use the base + # dataparser config. + DataParserUnion = DataParserConfig +else: + # At runtime, populate a Union type dynamically. This is used by `tyro` to generate + # subcommands in the CLI. + DataParserUnion = tyro.extras.subcommand_type_from_defaults( all_dataparsers, prefix_names=False, # Omit prefixes in subcommands themselves. ) -] -"""Union over possible dataparser types, annotated with metadata for tyro. This is the -same as the vanilla union, but results in shorter subcommand names.""" + +AnnotatedDataParserUnion = tyro.conf.OmitSubcommandPrefixes[DataParserUnion] # Omit prefixes of flags in subcommands. +"""Union over possible dataparser types, annotated with metadata for tyro. This is +the same as the vanilla union, but results in shorter subcommand names."""