Add file extensions to open and save dialogs (#532)

* Add file extensions to Save As dialog

* Include leading dots in file extensions for simplicity

* Add file extensions to open mappings dialogs

* Remove unused tinyMappingsFileChooser

* Use the same file chooser for all mapping IO

* Fix NPE by using Enigma directories as the default mapping format

Fixes #533.

* Fix code style

* Allow .mappings extension for single Enigma files

* gradlew.bat

---------

Co-authored-by: NebelNidas <nebelnidas@gmail.com>

6 files changed, 197 insertions, 35 deletions

diff --git a/enigma-swing/src/main/java/cuchaz/enigma/gui/Gui.java b/enigma-swing/src/main/java/cuchaz/enigma/gui/Gui.java
index 11013153..05146d4a 100644
--- a/enigma-swing/src/main/java/cuchaz/enigma/gui/Gui.java
+++ b/enigma-swing/src/main/java/cuchaz/enigma/gui/Gui.java
@@ -63,6 +63,7 @@ import cuchaz.enigma.gui.panels.IdentifierPanel;
 import cuchaz.enigma.gui.panels.ObfPanel;
 import cuchaz.enigma.gui.panels.StructurePanel;
 import cuchaz.enigma.gui.renderer.MessageListCellRenderer;
+import cuchaz.enigma.gui.util.ExtensionFileFilter;
 import cuchaz.enigma.gui.util.GuiUtil;
 import cuchaz.enigma.gui.util.LanguageUtil;
 import cuchaz.enigma.gui.util.ScaleUtil;
@@ -117,8 +118,7 @@ public class Gui {
         private final JLabel connectionStatusLabel = new JLabel();
 
         public final JFileChooser jarFileChooser = new JFileChooser();
-        public final JFileChooser tinyMappingsFileChooser = new JFileChooser();
-        public final JFileChooser enigmaMappingsFileChooser = new JFileChooser();
+        public final JFileChooser mappingsFileChooser = new JFileChooser();
         public final JFileChooser exportSourceFileChooser = new JFileChooser();
         public final JFileChooser exportJarFileChooser = new JFileChooser();
         public SearchDialog searchDialog;
@@ -147,10 +147,6 @@ public class Gui {
 
         private void setupUi() {
                 this.jarFileChooser.setFileSelectionMode(JFileChooser.FILES_ONLY);
-                this.tinyMappingsFileChooser.setFileSelectionMode(JFileChooser.FILES_ONLY);
-
-                this.enigmaMappingsFileChooser.setFileSelectionMode(JFileChooser.FILES_AND_DIRECTORIES);
-                this.enigmaMappingsFileChooser.setAcceptAllFileFilterUsed(false);
 
                 this.exportSourceFileChooser.setFileSelectionMode(JFileChooser.DIRECTORIES_ONLY);
                 this.exportSourceFileChooser.setAcceptAllFileFilterUsed(false);
@@ -322,7 +318,7 @@ public class Gui {
         }
 
         public void setMappingsFile(Path path) {
-                this.enigmaMappingsFileChooser.setSelectedFile(path != null ? path.toFile() : null);
+                this.mappingsFileChooser.setSelectedFile(path != null ? path.toFile() : null);
                 updateUiState();
         }
 
@@ -436,8 +432,10 @@ public class Gui {
         }
 
         public CompletableFuture<Void> saveMapping() {
-                if (this.enigmaMappingsFileChooser.getSelectedFile() != null || this.enigmaMappingsFileChooser.showSaveDialog(this.mainWindow.frame()) == JFileChooser.APPROVE_OPTION) {
-                        return this.controller.saveMappings(this.enigmaMappingsFileChooser.getSelectedFile().toPath());
+                ExtensionFileFilter.setupFileChooser(this.mappingsFileChooser, this.controller.getLoadedMappingFormat());
+
+                if (this.mappingsFileChooser.getSelectedFile() != null || this.mappingsFileChooser.showSaveDialog(this.mainWindow.frame()) == JFileChooser.APPROVE_OPTION) {
+                        return this.controller.saveMappings(ExtensionFileFilter.getSavePath(this.mappingsFileChooser));
                 }
 
                 return CompletableFuture.completedFuture(null);
diff --git a/enigma-swing/src/main/java/cuchaz/enigma/gui/GuiController.java b/enigma-swing/src/main/java/cuchaz/enigma/gui/GuiController.java
index c5799019..ad10abfc 100644
--- a/enigma-swing/src/main/java/cuchaz/enigma/gui/GuiController.java
+++ b/enigma-swing/src/main/java/cuchaz/enigma/gui/GuiController.java
@@ -99,7 +99,7 @@ public class GuiController implements ClientPacketHandler {
         private IndexTreeBuilder indexTreeBuilder;
 
         private Path loadedMappingPath;
-        private MappingFormat loadedMappingFormat;
+        private MappingFormat loadedMappingFormat = MappingFormat.ENIGMA_DIRECTORY;
 
         private ClassHandleProvider chp;
 
@@ -180,6 +180,10 @@ public class GuiController implements ClientPacketHandler {
                 chp.invalidateJavadoc();
         }
 
+        public MappingFormat getLoadedMappingFormat() {
+                return loadedMappingFormat;
+        }
+
         public CompletableFuture<Void> saveMappings(Path path) {
                 return saveMappings(path, loadedMappingFormat);
         }
diff --git a/enigma-swing/src/main/java/cuchaz/enigma/gui/elements/MenuBar.java b/enigma-swing/src/main/java/cuchaz/enigma/gui/elements/MenuBar.java
index 30e35861..35129799 100644
--- a/enigma-swing/src/main/java/cuchaz/enigma/gui/elements/MenuBar.java
+++ b/enigma-swing/src/main/java/cuchaz/enigma/gui/elements/MenuBar.java
@@ -34,6 +34,7 @@ import cuchaz.enigma.gui.dialog.CreateServerDialog;
 import cuchaz.enigma.gui.dialog.FontDialog;
 import cuchaz.enigma.gui.dialog.SearchDialog;
 import cuchaz.enigma.gui.dialog.StatsDialog;
+import cuchaz.enigma.gui.util.ExtensionFileFilter;
 import cuchaz.enigma.gui.util.GuiUtil;
 import cuchaz.enigma.gui.util.LanguageUtil;
 import cuchaz.enigma.gui.util.ScaleUtil;
@@ -175,7 +176,7 @@ public class MenuBar {
 
                 this.jarCloseItem.setEnabled(jarOpen);
                 this.openMappingsMenu.setEnabled(jarOpen);
-                this.saveMappingsItem.setEnabled(jarOpen && this.gui.enigmaMappingsFileChooser.getSelectedFile() != null && connectionState != ConnectionState.CONNECTED);
+                this.saveMappingsItem.setEnabled(jarOpen && this.gui.mappingsFileChooser.getSelectedFile() != null && connectionState != ConnectionState.CONNECTED);
                 this.saveMappingsAsMenu.setEnabled(jarOpen);
                 this.closeMappingsItem.setEnabled(jarOpen);
                 this.reloadMappingsItem.setEnabled(jarOpen);
@@ -250,7 +251,7 @@ public class MenuBar {
         }
 
         private void onSaveMappingsClicked() {
-                this.gui.getController().saveMappings(this.gui.enigmaMappingsFileChooser.getSelectedFile().toPath());
+                this.gui.getController().saveMappings(this.gui.mappingsFileChooser.getSelectedFile().toPath());
         }
 
         private void openMappingsDiscardPrompt(Runnable then) {
@@ -422,12 +423,13 @@ public class MenuBar {
         private static void addOpenMappingsMenuEntry(String text, MappingFormat format, boolean mappingIo, JMenu openMappingsMenu, Gui gui) {
                 JMenuItem item = new JMenuItem(text);
                 item.addActionListener(event -> {
-                        gui.enigmaMappingsFileChooser.setCurrentDirectory(new File(UiConfig.getLastSelectedDir()));
+                        ExtensionFileFilter.setupFileChooser(gui.mappingsFileChooser, format);
+                        gui.mappingsFileChooser.setCurrentDirectory(new File(UiConfig.getLastSelectedDir()));
 
-                        if (gui.enigmaMappingsFileChooser.showOpenDialog(gui.getFrame()) == JFileChooser.APPROVE_OPTION) {
-                                File selectedFile = gui.enigmaMappingsFileChooser.getSelectedFile();
+                        if (gui.mappingsFileChooser.showOpenDialog(gui.getFrame()) == JFileChooser.APPROVE_OPTION) {
+                                File selectedFile = gui.mappingsFileChooser.getSelectedFile();
                                 gui.getController().openMappings(format, selectedFile.toPath(), mappingIo);
-                                UiConfig.setLastSelectedDir(gui.enigmaMappingsFileChooser.getCurrentDirectory().toString());
+                                UiConfig.setLastSelectedDir(gui.mappingsFileChooser.getCurrentDirectory().toString());
                         }
                 });
                 openMappingsMenu.add(item);
@@ -453,15 +455,18 @@ public class MenuBar {
         private static void addSaveMappingsAsMenuEntry(String text, MappingFormat format, boolean mappingIo, JMenu saveMappingsAsMenu, JMenuItem saveMappingsItem, Gui gui) {
                 JMenuItem item = new JMenuItem(text);
                 item.addActionListener(event -> {
-                        // TODO: Use a specific file chooser for it
-                        if (gui.enigmaMappingsFileChooser.getCurrentDirectory() == null) {
-                                gui.enigmaMappingsFileChooser.setCurrentDirectory(new File(UiConfig.getLastSelectedDir()));
+                        JFileChooser fileChooser = gui.mappingsFileChooser;
+                        ExtensionFileFilter.setupFileChooser(fileChooser, format);
+
+                        if (fileChooser.getCurrentDirectory() == null) {
+                                fileChooser.setCurrentDirectory(new File(UiConfig.getLastSelectedDir()));
                         }
 
-                        if (gui.enigmaMappingsFileChooser.showSaveDialog(gui.getFrame()) == JFileChooser.APPROVE_OPTION) {
-                                gui.getController().saveMappings(gui.enigmaMappingsFileChooser.getSelectedFile().toPath(), format, mappingIo);
+                        if (fileChooser.showSaveDialog(gui.getFrame()) == JFileChooser.APPROVE_OPTION) {
+                                Path savePath = ExtensionFileFilter.getSavePath(fileChooser);
+                                gui.getController().saveMappings(savePath, format, mappingIo);
                                 saveMappingsItem.setEnabled(true);
-                                UiConfig.setLastSelectedDir(gui.enigmaMappingsFileChooser.getCurrentDirectory().toString());
+                                UiConfig.setLastSelectedDir(fileChooser.getCurrentDirectory().toString());
                         }
                 });
                 saveMappingsAsMenu.add(item);
diff --git a/enigma-swing/src/main/java/cuchaz/enigma/gui/util/ExtensionFileFilter.java b/enigma-swing/src/main/java/cuchaz/enigma/gui/util/ExtensionFileFilter.java
new file mode 100644
index 00000000..a514cae5
--- /dev/null
+++ b/enigma-swing/src/main/java/cuchaz/enigma/gui/util/ExtensionFileFilter.java
@@ -0,0 +1,118 @@
+package cuchaz.enigma.gui.util;
+
+import java.io.File;
+import java.nio.file.Path;
+import java.util.List;
+import java.util.Locale;
+import java.util.StringJoiner;
+
+import javax.swing.JFileChooser;
+import javax.swing.filechooser.FileFilter;
+
+import cuchaz.enigma.translation.mapping.serde.MappingFormat;
+import cuchaz.enigma.utils.I18n;
+
+public final class ExtensionFileFilter extends FileFilter {
+        private final String formatName;
+        private final List<String> extensions;
+
+        /**
+         * Constructs an {@code ExtensionFileFilter}.
+         *
+         * @param formatName the human-readable name of the file format
+         * @param extensions the file extensions with their leading dots (e.g. {@code .txt})
+         */
+        public ExtensionFileFilter(String formatName, List<String> extensions) {
+                this.formatName = formatName;
+                this.extensions = extensions;
+        }
+
+        public List<String> getExtensions() {
+                return extensions;
+        }
+
+        @Override
+        public boolean accept(File f) {
+                // Always accept directories so the user can see them.
+                if (f.isDirectory()) {
+                        return true;
+                }
+
+                for (String extension : extensions) {
+                        if (f.getName().endsWith(extension)) {
+                                return true;
+                        }
+                }
+
+                return false;
+        }
+
+        @Override
+        public String getDescription() {
+                var joiner = new StringJoiner(", ");
+
+                for (String extension : extensions) {
+                        joiner.add("*" + extension);
+                }
+
+                return I18n.translateFormatted("menu.file.mappings.file_filter", formatName, joiner.toString());
+        }
+
+        /**
+         * Sets up a file chooser with a mapping format. This method resets the choosable filters,
+         * and adds and selects a new filter based on the provided mapping format.
+         *
+         * @param fileChooser the mapping format
+         */
+        public static void setupFileChooser(JFileChooser fileChooser, MappingFormat format) {
+                // Remove previous custom filters.
+                fileChooser.resetChoosableFileFilters();
+
+                if (format.getFileType().isDirectory()) {
+                        fileChooser.setFileSelectionMode(JFileChooser.DIRECTORIES_ONLY);
+                } else {
+                        fileChooser.setFileSelectionMode(JFileChooser.FILES_ONLY);
+                        String formatName = I18n.translate("mapping_format." + format.name().toLowerCase(Locale.ROOT));
+                        var filter = new ExtensionFileFilter(formatName, format.getFileType().extensions());
+                        // Add our new filter to the list...
+                        fileChooser.addChoosableFileFilter(filter);
+                        // ...and choose it as the default.
+                        fileChooser.setFileFilter(filter);
+                }
+        }
+
+        /**
+         * Fixes a missing file extension in a save file path when the selected filter
+         * is an {@code ExtensionFileFilter}.
+         *
+         * @param fileChooser the file chooser to check
+         * @return the fixed path
+         */
+        public static Path getSavePath(JFileChooser fileChooser) {
+                Path savePath = fileChooser.getSelectedFile().toPath();
+
+                if (fileChooser.getFileFilter() instanceof ExtensionFileFilter extensionFilter) {
+                        // Check that the file name ends with the extension.
+                        String fileName = savePath.getFileName().toString();
+                        boolean hasExtension = false;
+
+                        for (String extension : extensionFilter.getExtensions()) {
+                                if (fileName.endsWith(extension)) {
+                                        hasExtension = true;
+                                        break;
+                                }
+                        }
+
+                        if (!hasExtension) {
+                                String defaultExtension = extensionFilter.getExtensions().get(0);
+                                // If not, add the extension.
+                                savePath = savePath.resolveSibling(fileName + defaultExtension);
+                                // Store the adjusted file, so that it shows up properly
+                                // the next time this dialog is used.
+                                fileChooser.setSelectedFile(savePath.toFile());
+                        }
+                }
+
+                return savePath;
+        }
+}
diff --git a/enigma/src/main/java/cuchaz/enigma/translation/mapping/serde/MappingFormat.java b/enigma/src/main/java/cuchaz/enigma/translation/mapping/serde/MappingFormat.java
index 4521591e..aeac9db4 100644
--- a/enigma/src/main/java/cuchaz/enigma/translation/mapping/serde/MappingFormat.java
+++ b/enigma/src/main/java/cuchaz/enigma/translation/mapping/serde/MappingFormat.java
@@ -33,28 +33,30 @@ import cuchaz.enigma.translation.mapping.tree.EntryTree;
 import cuchaz.enigma.utils.I18n;
 
 public enum MappingFormat {
-        ENIGMA_FILE(EnigmaMappingsWriter.FILE, EnigmaMappingsReader.FILE, net.fabricmc.mappingio.format.MappingFormat.ENIGMA_FILE, true),
-        ENIGMA_DIRECTORY(EnigmaMappingsWriter.DIRECTORY, EnigmaMappingsReader.DIRECTORY, net.fabricmc.mappingio.format.MappingFormat.ENIGMA_DIR, true),
-        ENIGMA_ZIP(EnigmaMappingsWriter.ZIP, EnigmaMappingsReader.ZIP, null, false),
-        TINY_V2(new TinyV2Writer("intermediary", "named"), new TinyV2Reader(), net.fabricmc.mappingio.format.MappingFormat.TINY_2_FILE, true),
-        TINY_FILE(TinyMappingsWriter.INSTANCE, TinyMappingsReader.INSTANCE, net.fabricmc.mappingio.format.MappingFormat.TINY_FILE, true),
-        SRG_FILE(SrgMappingsWriter.INSTANCE, null, net.fabricmc.mappingio.format.MappingFormat.SRG_FILE, true),
-        XSRG_FILE(null, null, net.fabricmc.mappingio.format.MappingFormat.XSRG_FILE, true),
-        CSRG_FILE(null, null, net.fabricmc.mappingio.format.MappingFormat.CSRG_FILE, false),
-        TSRG_FILE(null, null, net.fabricmc.mappingio.format.MappingFormat.TSRG_FILE, false),
-        TSRG_2_FILE(null, null, net.fabricmc.mappingio.format.MappingFormat.TSRG_2_FILE, false),
-        PROGUARD(null, ProguardMappingsReader.INSTANCE, net.fabricmc.mappingio.format.MappingFormat.PROGUARD_FILE, true),
-        RECAF(RecafMappingsWriter.INSTANCE, RecafMappingsReader.INSTANCE, null, false);
+        ENIGMA_FILE(EnigmaMappingsWriter.FILE, EnigmaMappingsReader.FILE, FileType.MAPPING, net.fabricmc.mappingio.format.MappingFormat.ENIGMA_FILE, true),
+        ENIGMA_DIRECTORY(EnigmaMappingsWriter.DIRECTORY, EnigmaMappingsReader.DIRECTORY, FileType.DIRECTORY, net.fabricmc.mappingio.format.MappingFormat.ENIGMA_DIR, true),
+        ENIGMA_ZIP(EnigmaMappingsWriter.ZIP, EnigmaMappingsReader.ZIP, FileType.ZIP, null, false),
+        TINY_V2(new TinyV2Writer("intermediary", "named"), new TinyV2Reader(), FileType.TINY, net.fabricmc.mappingio.format.MappingFormat.TINY_2_FILE, true),
+        TINY_FILE(TinyMappingsWriter.INSTANCE, TinyMappingsReader.INSTANCE, FileType.TINY, net.fabricmc.mappingio.format.MappingFormat.TINY_FILE, true),
+        SRG_FILE(SrgMappingsWriter.INSTANCE, null, FileType.SRG, net.fabricmc.mappingio.format.MappingFormat.SRG_FILE, true),
+        XSRG_FILE(null, null, FileType.XSRG, net.fabricmc.mappingio.format.MappingFormat.XSRG_FILE, true),
+        CSRG_FILE(null, null, FileType.CSRG, net.fabricmc.mappingio.format.MappingFormat.CSRG_FILE, false),
+        TSRG_FILE(null, null, FileType.TSRG, net.fabricmc.mappingio.format.MappingFormat.TSRG_FILE, false),
+        TSRG_2_FILE(null, null, FileType.TSRG, net.fabricmc.mappingio.format.MappingFormat.TSRG_2_FILE, false),
+        PROGUARD(null, ProguardMappingsReader.INSTANCE, FileType.TXT, net.fabricmc.mappingio.format.MappingFormat.PROGUARD_FILE, true),
+        RECAF(RecafMappingsWriter.INSTANCE, RecafMappingsReader.INSTANCE, FileType.TXT, null, false);
 
         private final MappingsWriter writer;
         private final MappingsReader reader;
+        private final FileType fileType;
         private final net.fabricmc.mappingio.format.MappingFormat mappingIoCounterpart;
         private final boolean hasMappingIoWriter;
         private boolean lastUsedMappingIoWriter;
 
-        MappingFormat(MappingsWriter writer, MappingsReader reader, net.fabricmc.mappingio.format.MappingFormat mappingIoCounterpart, boolean hasMappingIoWriter) {
+        MappingFormat(MappingsWriter writer, MappingsReader reader, FileType fileType, net.fabricmc.mappingio.format.MappingFormat mappingIoCounterpart, boolean hasMappingIoWriter) {
                 this.writer = writer;
                 this.reader = reader;
+                this.fileType = fileType;
                 this.mappingIoCounterpart = mappingIoCounterpart;
                 this.hasMappingIoWriter = hasMappingIoWriter;
         }
@@ -126,6 +128,11 @@ public enum MappingFormat {
                 return reader;
         }
 
+        @ApiStatus.Internal
+        public FileType getFileType() {
+                return fileType;
+        }
+
         @Nullable
         @ApiStatus.Internal
         public net.fabricmc.mappingio.format.MappingFormat getMappingIoCounterpart() {
@@ -164,4 +171,33 @@ public enum MappingFormat {
                                 .filter(MappingFormat::isWritable)
                                 .toList();
         }
+
+        /**
+         * A file type. It can be either a single file with an extension, or a directory
+         * with a {@code null} extension.
+         *
+         * <p>If a file type has multiple extensions, the default for saving will be the first one.
+         *
+         * @param extensions the file extensions with the leading dot {@code .}, or an empty list for a directory
+         */
+        @ApiStatus.Internal
+        public record FileType(List<String> extensions) {
+                public static final FileType DIRECTORY = new FileType();
+                public static final FileType MAPPING = new FileType(".mapping", ".mappings");
+                public static final FileType SRG = new FileType(".srg");
+                public static final FileType XSRG = new FileType(".xsrg");
+                public static final FileType CSRG = new FileType(".csrg");
+                public static final FileType TSRG = new FileType(".tsrg");
+                public static final FileType TINY = new FileType(".tiny");
+                public static final FileType TXT = new FileType(".txt");
+                public static final FileType ZIP = new FileType(".zip");
+
+                public FileType(String... extensions) {
+                        this(List.of(extensions));
+                }
+
+                public boolean isDirectory() {
+                        return extensions.isEmpty();
+                }
+        }
 }
diff --git a/enigma/src/main/resources/lang/en_us.json b/enigma/src/main/resources/lang/en_us.json
index 2a3cfb2d..42228144 100644
--- a/enigma/src/main/resources/lang/en_us.json
+++ b/enigma/src/main/resources/lang/en_us.json
@@ -22,6 +22,7 @@
         "menu.file.mappings.save_as": "Save Mappings As...",
         "menu.file.mappings.close": "Close Mappings",
         "menu.file.mappings.drop": "Drop Invalid Mappings",
+        "menu.file.mappings.file_filter": "%s (%s)",
         "menu.file.reload_mappings": "Reload Mappings",
         "menu.file.reload_all": "Reload Jar/Mappings",
         "menu.file.export.source": "Export Source...",