Merge pull request #41 from RatzzFatzz/dev

Improve subtitle recognition

src/main/java/at/pcgamingfreaks/mkvaudiosubtitlechanger/config/Config.java

@@ -42,6 +42,7 @@ public class Config {
     private Set<String> forcedKeywords = new HashSet<>(Arrays.asList("forced", "signs", "songs"));
     private Set<String> commentaryKeywords = new HashSet<>(Arrays.asList("commentary", "director"));
     private Set<String> excludedDirectories = new HashSet<>();
+    private Set<String> preferredSubtitles = new HashSet<>(Arrays.asList("unstyled"));
 
     private List<AttributeConfig> attributeConfig;
 

src/main/java/at/pcgamingfreaks/mkvaudiosubtitlechanger/config/ConfigLoader.java

@@ -28,6 +28,7 @@ public class ConfigLoader {
             new SetValidator(FORCED_KEYWORDS, false, true),
             new SetValidator(COMMENTARY_KEYWORDS, false, true),
             new SetValidator(EXCLUDED_DIRECTORY, false, true),
+            new SetValidator(PREFERRED_SUBTITLES, false, true),
             new AttributeConfigValidator(),
             new CoherentConfigValidator(COHERENT, false),
             new BooleanValidator(FORCE_COHERENT, false)

src/main/java/at/pcgamingfreaks/mkvaudiosubtitlechanger/config/validator/ConfigValidator.java

@@ -142,7 +142,7 @@ public abstract class ConfigValidator<FieldType> {
 
 	protected Predicate<Method> containsGetterOf(ConfigProperty property) {
 		return method -> StringUtils.startsWith(method.getName(), "get")
-				&& StringUtils.containsIgnoreCase(method.getName(), property.prop().replace("-", ""));
+				&& StringUtils.equalsIgnoreCase(method.getName().replace("get", ""), property.prop().replace("-", ""));
 	}
 
 

src/main/java/at/pcgamingfreaks/mkvaudiosubtitlechanger/impl/MkvFileProcessor.java

@@ -15,6 +15,7 @@ import java.io.InputStream;
 import java.io.InputStreamReader;
 import java.util.*;
 import java.util.stream.Collectors;
+import java.util.stream.Stream;
 
 import static at.pcgamingfreaks.mkvaudiosubtitlechanger.model.LaneType.AUDIO;
 import static at.pcgamingfreaks.mkvaudiosubtitlechanger.model.LaneType.SUBTITLES;
@@ -23,6 +24,10 @@ import static java.lang.String.format;
 @Slf4j
 public class MkvFileProcessor implements FileProcessor {
     private final ObjectMapper mapper = new ObjectMapper();
+
+    private static final SubtitleTrackComparator subtitleTrackComparator =
+            new SubtitleTrackComparator(Config.getInstance().getPreferredSubtitles().toArray(new String[0]));
+
     private static final String DISABLE_DEFAULT_TRACK = "--edit track:%s --set flag-default=0 ";
     private static final String ENABLE_DEFAULT_TRACK = "--edit track:%s --set flag-default=1 ";
     private static final String ENABLE_FORCED_TRACK = "--edit track:%s --set flag-forced=1 ";
@@ -98,23 +103,31 @@ public class MkvFileProcessor implements FileProcessor {
     @Override
     public void detectDesiredTracks(FileInfoDto info, List<FileAttribute> nonForcedTracks, List<FileAttribute> nonCommentaryTracks,
                                     AttributeConfig... configs) {
+        Set<FileAttribute> tracks = SetUtils.retainOf(nonForcedTracks, nonCommentaryTracks);
+        Set<FileAttribute> audioTracks = tracks.stream().filter(a -> AUDIO.equals(a.getType())).collect(Collectors.toSet());
+        Set<FileAttribute> subtitleTracks = tracks.stream().filter(a -> SUBTITLES.equals(a.getType())).collect(Collectors.toSet());
+
         for (AttributeConfig config : configs) {
-            FileAttribute desiredAudio = null;
+            Optional<FileAttribute> desiredAudio = detectDesiredTrack(config.getAudioLanguage(), audioTracks).findFirst();
-            FileAttribute desiredSubtitle = null;
+            Optional<FileAttribute> desiredSubtitle = detectDesiredSubtitleTrack(config.getSubtitleLanguage(), subtitleTracks).findFirst();
-            for (FileAttribute attribute : SetUtils.retainOf(nonForcedTracks, nonCommentaryTracks)) {
+
-                if (attribute.getLanguage().equals(config.getAudioLanguage())
+            if (desiredAudio.isPresent() && desiredSubtitle.isPresent()) {
-                        && AUDIO.equals(attribute.getType())) desiredAudio = attribute;
+                info.setDesiredAudioLane(desiredAudio.get());
-                if (attribute.getLanguage().equals(config.getSubtitleLanguage())
+                info.setDesiredSubtitleLane(desiredSubtitle.get());
-                        && SUBTITLES.equals(attribute.getType())) desiredSubtitle = attribute;
-            }
-            if (desiredAudio != null && desiredSubtitle != null) {
-                info.setDesiredAudioLane(desiredAudio);
-                info.setDesiredSubtitleLane(desiredSubtitle);
                 break;
             }
         }
     }
 
+    private Stream<FileAttribute> detectDesiredTrack(String language, Set<FileAttribute> tracks) {
+        return tracks.stream().filter(track -> language.equals(track.getLanguage()));
+    }
+
+    private Stream<FileAttribute> detectDesiredSubtitleTrack(String language, Set<FileAttribute> tracks) {
+        return detectDesiredTrack(language, tracks)
+                .sorted(subtitleTrackComparator.reversed());
+    }
+
     @Override
     public List<FileAttribute> retrieveNonForcedTracks(List<FileAttribute> attributes) {
         return attributes.stream()

src/main/java/at/pcgamingfreaks/mkvaudiosubtitlechanger/impl/SubtitleTrackComparator.java

@@ -0,0 +1,34 @@
+package at.pcgamingfreaks.mkvaudiosubtitlechanger.impl;
+
+import at.pcgamingfreaks.mkvaudiosubtitlechanger.model.FileAttribute;
+import lombok.RequiredArgsConstructor;
+import org.apache.commons.lang3.StringUtils;
+
+import java.util.Comparator;
+
+@RequiredArgsConstructor
+public class SubtitleTrackComparator implements Comparator<FileAttribute> {
+    private final String[] preferredSubtitles;
+
+    /**
+     * {@inheritDoc}
+     */
+    @Override
+    public int compare(FileAttribute track1, FileAttribute track2) {
+        int result = 0;
+
+        if (StringUtils.containsAnyIgnoreCase(track1.getTrackName(), preferredSubtitles)) {
+            result++;
+        }
+        if (StringUtils.containsAnyIgnoreCase(track2.getTrackName(), preferredSubtitles)) {
+            result--;
+        }
+
+        if (result == 0) {
+            if (track1.isDefaultTrack()) result++;
+            if (track2.isDefaultTrack()) result--;
+        }
+
+        return result;
+    }
+}

src/main/java/at/pcgamingfreaks/mkvaudiosubtitlechanger/impl/kernel/AttributeUpdaterKernel.java

@@ -4,6 +4,7 @@ import at.pcgamingfreaks.mkvaudiosubtitlechanger.config.Config;
 import at.pcgamingfreaks.mkvaudiosubtitlechanger.exceptions.MkvToolNixException;
 import at.pcgamingfreaks.mkvaudiosubtitlechanger.impl.FileCollector;
 import at.pcgamingfreaks.mkvaudiosubtitlechanger.impl.FileProcessor;
+import at.pcgamingfreaks.mkvaudiosubtitlechanger.model.AttributeConfig;
 import at.pcgamingfreaks.mkvaudiosubtitlechanger.model.FileAttribute;
 import at.pcgamingfreaks.mkvaudiosubtitlechanger.model.FileInfoDto;
 import at.pcgamingfreaks.mkvaudiosubtitlechanger.model.ResultStatistic;
@@ -93,7 +94,8 @@ public abstract class AttributeUpdaterKernel {
         List<FileAttribute> nonCommentaryTracks = processor.retrieveNonCommentaryTracks(attributes);
 
         processor.detectDefaultTracks(fileInfo, attributes, nonForcedTracks);
-        processor.detectDesiredTracks(fileInfo, nonForcedTracks, nonCommentaryTracks);
+        processor.detectDesiredTracks(fileInfo, nonForcedTracks, nonCommentaryTracks,
+                Config.getInstance().getAttributeConfig().toArray(new AttributeConfig[]{}));
 
         updateFile(fileInfo);
     }

src/main/java/at/pcgamingfreaks/mkvaudiosubtitlechanger/model/ConfigProperty.java

@@ -23,10 +23,26 @@ public enum ConfigProperty {
     EXCLUDED_DIRECTORY("excluded-directories", "Directories to be excluded, combines with config file", "e", Option.UNLIMITED_VALUES),
     FORCED_KEYWORDS("forced-keywords", "Additional keywords to identify forced tracks", "fk", Option.UNLIMITED_VALUES),
     COMMENTARY_KEYWORDS("commentary-keywords", "Additional keywords to identify commentary tracks", "ck", Option.UNLIMITED_VALUES),
+    PREFERRED_SUBTITLES("preferred-subtitles", "Additional keywords to prefer specific subtitle tracks", "ps", Option.UNLIMITED_VALUES),
     ARGUMENTS("arguments", "List of arguments", null, 0),
     VERSION("version", "Display version", "v", 0),
     HELP("help", "\"For help this is\" - Yoda", "h", 0);
 
+    /*
+     * Verify at startup that there are no duplicated shortParameters.
+     */
+    static {
+        Set<String> shortParameters = new HashSet<>();
+        for (String param : Arrays.stream(ConfigProperty.values()).map(ConfigProperty::abrv).collect(Collectors.toList())) {
+            if (shortParameters.contains(param)) {
+                throw new IllegalStateException("It is not allowed to have multiple properties with the same abbreviation!");
+            }
+            if (param != null) {
+                shortParameters.add(param);
+            }
+        }
+    }
+
     private final String property;
     private final String description;
     private final String shortParameter;
@@ -47,19 +63,4 @@ public enum ConfigProperty {
     public int args() {
         return args;
     }
-
-    /*
-     * Verify at startup that there are no duplicated shortParameters.
-     */
-    static {
-        Set<String> shortParameters = new HashSet<>();
-        for (String param: Arrays.stream(ConfigProperty.values()).map(ConfigProperty::abrv).collect(Collectors.toList())) {
-            if (shortParameters.contains(param)) {
-                throw new IllegalStateException("It is not allowed to have multiple properties with the same abbreviation!");
-            }
-            if (param != null) {
-                shortParameters.add(param);
-            }
-        }
-    }
 }

src/main/java/at/pcgamingfreaks/mkvaudiosubtitlechanger/model/FileAttribute.java

@@ -4,6 +4,8 @@ import lombok.AllArgsConstructor;
 import lombok.Getter;
 import lombok.extern.slf4j.Slf4j;
 
+import java.util.Objects;
+
 @Slf4j
 @Getter
 @AllArgsConstructor
@@ -15,6 +17,24 @@ public class FileAttribute {
     private final boolean forcedTrack;
     private final LaneType type;
 
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) return true;
+        if (o == null || getClass() != o.getClass()) return false;
+        FileAttribute attribute = (FileAttribute) o;
+        return id == attribute.id
+                && defaultTrack == attribute.defaultTrack
+                && forcedTrack == attribute.forcedTrack
+                && Objects.equals(language, attribute.language)
+                && Objects.equals(trackName, attribute.trackName)
+                && type == attribute.type;
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(id, language, trackName, defaultTrack, forcedTrack, type);
+    }
+
     @Override
     public String toString() {
         return "[" + "id=" + id +

src/test/java/at/pcgamingfreaks/mkvaudiosubtitlechanger/impl/SubtitleTrackComparatorTest.java

@@ -0,0 +1,46 @@
+package at.pcgamingfreaks.mkvaudiosubtitlechanger.impl;
+
+import at.pcgamingfreaks.mkvaudiosubtitlechanger.model.FileAttribute;
+import at.pcgamingfreaks.mkvaudiosubtitlechanger.model.LaneType;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.Arguments;
+import org.junit.jupiter.params.provider.MethodSource;
+
+import java.util.List;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
+
+import static org.junit.jupiter.api.Assertions.*;
+
+class SubtitleTrackComparatorTest {
+    private static final SubtitleTrackComparator comparator = new SubtitleTrackComparator(new String[]{"unstyled"});
+
+    private static Stream<Arguments> compareArguments() {
+        return Stream.of(
+                Arguments.of(List.of(attr("unstyled sub", false), attr("styled sub", false)),
+                        List.of(attr("unstyled sub", false), attr("styled sub", false))),
+                Arguments.of(List.of(attr("styled sub", false), attr("unstyled sub", false)),
+                        List.of(attr("unstyled sub", false), attr("styled sub", false))),
+
+                Arguments.of(List.of(attr("unstyled sub", true), attr("styled sub", false)),
+                        List.of(attr("unstyled sub", true), attr("styled sub", false))),
+                Arguments.of(List.of(attr("styled sub", true), attr("unstyled sub", false)),
+                        List.of(attr("unstyled sub", false), attr("styled sub", true))),
+
+                Arguments.of(List.of(attr("unstyled sub", true), attr("unstyled sub", false)),
+                        List.of(attr("unstyled sub", true), attr("unstyled sub", false)))
+        );
+    }
+
+    @ParameterizedTest
+    @MethodSource("compareArguments")
+    void compare(List<FileAttribute> input, List<FileAttribute> expected) {
+        List<FileAttribute> result = input.stream().sorted(comparator.reversed()).collect(Collectors.toList());
+
+        assertArrayEquals(expected.toArray(new FileAttribute[0]), result.toArray(new FileAttribute[0]));
+    }
+
+    private static FileAttribute attr(String trackName, boolean defaultTrack) {
+        return new FileAttribute(0, "", trackName, defaultTrack, false, LaneType.SUBTITLES);
+    }
+}