diff --git a/io_scene_psk_psa/psa/config.py b/io_scene_psk_psa/psa/config.py index f9e1515..184d59c 100644 --- a/io_scene_psk_psa/psa/config.py +++ b/io_scene_psk_psa/psa/config.py @@ -13,36 +13,66 @@ class PsaConfig: self.sequence_bone_flags: Dict[str, Dict[int, int]] = dict() +def _load_config_file(file_path: str) -> ConfigParser: + """ + UEViewer exports a dialect of INI files that is not compatible with Python's ConfigParser. + Specifically, it allows values in this format: + + [Section] + Key1 + Key2 + + This is not allowed in Python's ConfigParser, which requires a '=' character after each key name. + To work around this, we'll modify the file to add the '=' character after each key name if it is missing. + """ + with open(file_path, 'r') as f: + lines = f.read().split('\n') + + lines = [re.sub(r'^\s*(\w+)\s*$', r'\1=', line) for line in lines] + + contents = '\n'.join(lines) + + config = ConfigParser() + config.read_string(contents) + + return config + + +def _get_bone_flags_from_value(value: str) -> int: + match value: + case 'all': + return (REMOVE_TRACK_LOCATION | REMOVE_TRACK_ROTATION) + case 'trans': + return REMOVE_TRACK_LOCATION + case 'rot': + return REMOVE_TRACK_ROTATION + case _: + return 0 + + def read_psa_config(psa_reader: PsaReader, file_path: str) -> PsaConfig: psa_config = PsaConfig() - config = ConfigParser() - config.read(file_path) - - psa_sequence_names = list(psa_reader.sequences.keys()) - lowercase_sequence_names = [sequence_name.lower() for sequence_name in psa_sequence_names] + config = _load_config_file(file_path) if config.has_section('RemoveTracks'): for key, value in config.items('RemoveTracks'): match = re.match(f'^(.+)\.(\d+)$', key) sequence_name = match.group(1) - bone_index = int(match.group(2)) # Map the sequence name onto the actual sequence name in the PSA file. try: + psa_sequence_names = list(psa_reader.sequences.keys()) + lowercase_sequence_names = [sequence_name.lower() for sequence_name in psa_sequence_names] sequence_name = psa_sequence_names[lowercase_sequence_names.index(sequence_name.lower())] except ValueError: - pass + # Sequence name is not in the PSA file. + continue if sequence_name not in psa_config.sequence_bone_flags: psa_config.sequence_bone_flags[sequence_name] = dict() - match value: - case 'all': - psa_config.sequence_bone_flags[sequence_name][bone_index] = (REMOVE_TRACK_LOCATION | REMOVE_TRACK_ROTATION) - case 'trans': - psa_config.sequence_bone_flags[sequence_name][bone_index] = REMOVE_TRACK_LOCATION - case 'rot': - psa_config.sequence_bone_flags[sequence_name][bone_index] = REMOVE_TRACK_ROTATION + bone_index = int(match.group(2)) + psa_config.sequence_bone_flags[sequence_name][bone_index] = _get_bone_flags_from_value(value) return psa_config diff --git a/io_scene_psk_psa/psa/import_/operators.py b/io_scene_psk_psa/psa/import_/operators.py index afa964d..4b75a1a 100644 --- a/io_scene_psk_psa/psa/import_/operators.py +++ b/io_scene_psk_psa/psa/import_/operators.py @@ -158,6 +158,10 @@ class PSA_OT_import(Operator, ImportHelper): psa_reader = PsaReader(self.filepath) sequence_names = [x.action_name for x in pg.sequence_list if x.is_selected] + if len(sequence_names) == 0: + self.report({'ERROR_INVALID_CONTEXT'}, 'No sequences selected') + return {'CANCELLED'} + options = PsaImportOptions() options.sequence_names = sequence_names options.should_use_fake_user = pg.should_use_fake_user @@ -171,14 +175,14 @@ class PSA_OT_import(Operator, ImportHelper): options.fps_source = pg.fps_source options.fps_custom = pg.fps_custom - # Read the PSA config file if it exists. - config_path = Path(self.filepath).with_suffix('.config') - if config_path.exists(): - options.psa_config = read_psa_config(psa_reader, str(config_path)) - - if len(sequence_names) == 0: - self.report({'ERROR_INVALID_CONTEXT'}, 'No sequences selected') - return {'CANCELLED'} + if options.should_use_config_file: + # Read the PSA config file if it exists. + config_path = Path(self.filepath).with_suffix('.config') + if config_path.exists(): + try: + options.psa_config = read_psa_config(psa_reader, str(config_path)) + except Exception as e: + self.report({'WARNING'}, f'Failed to read PSA config file: {e}') result = import_psa(context, psa_reader, context.view_layer.objects.active, options) @@ -258,6 +262,8 @@ class PSA_OT_import(Operator, ImportHelper): col.use_property_decorate = False col.prop(pg, 'should_use_fake_user') col.prop(pg, 'should_stash') + col.prop(pg, 'should_use_config_file') + col.prop(pg, 'should_use_action_name_prefix') if pg.should_use_action_name_prefix: diff --git a/io_scene_psk_psa/psa/import_/properties.py b/io_scene_psk_psa/psa/import_/properties.py index b7e14f5..43dd375 100644 --- a/io_scene_psk_psa/psa/import_/properties.py +++ b/io_scene_psk_psa/psa/import_/properties.py @@ -32,6 +32,12 @@ class PSA_PG_import(PropertyGroup): description='Assign each imported action a fake user so that the data block is ' 'saved even it has no users', options=empty_set) + should_use_config_file: BoolProperty(default=True, name='Use Config File', + description='Use the .config file that is sometimes generated when the PSA ' + 'file is exported from UEViewer. This file contains ' + 'options that can be used to filter out certain bones tracks ' + 'from the imported actions', + options=empty_set) should_stash: BoolProperty(default=False, name='Stash', description='Stash each imported action as a strip on a new non-contributing NLA track', options=empty_set) diff --git a/io_scene_psk_psa/psa/importer.py b/io_scene_psk_psa/psa/importer.py index c4a66cc..5864645 100644 --- a/io_scene_psk_psa/psa/importer.py +++ b/io_scene_psk_psa/psa/importer.py @@ -24,6 +24,7 @@ class PsaImportOptions(object): self.bone_mapping_mode = 'CASE_INSENSITIVE' self.fps_source = 'SEQUENCE' self.fps_custom: float = 30.0 + self.should_use_config_file = True self.psa_config: PsaConfig = PsaConfig()