3 files changed, 26 insertions, 19 deletions

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 15afeed..ed0a2bf 100644
--- a/enigma-swing/src/main/java/cuchaz/enigma/gui/GuiController.java
+++ b/enigma-swing/src/main/java/cuchaz/enigma/gui/GuiController.java
@@ -484,9 +484,9 @@ public class GuiController implements ClientPacketHandler {
                 chp.invalidateMapped();
         }
 
-        public void openStats(Set<StatsMember> includedMembers) {
+        public void openStats(Set<StatsMember> includedMembers, String topLevelPackage) {
                 ProgressDialog.runOffThread(gui.getFrame(), progress -> {
-                        String data = new StatsGenerator(project).generate(progress, includedMembers);
+                        String data = new StatsGenerator(project).generate(progress, includedMembers, topLevelPackage);
 
                         try {
                                 File statsFile = File.createTempFile("stats", ".html");
diff --git a/enigma-swing/src/main/java/cuchaz/enigma/gui/dialog/StatsDialog.java b/enigma-swing/src/main/java/cuchaz/enigma/gui/dialog/StatsDialog.java
index 868eba7..d8d3acd 100644
--- a/enigma-swing/src/main/java/cuchaz/enigma/gui/dialog/StatsDialog.java
+++ b/enigma-swing/src/main/java/cuchaz/enigma/gui/dialog/StatsDialog.java
@@ -1,16 +1,13 @@
 package cuchaz.enigma.gui.dialog;
 
-import java.awt.BorderLayout;
+import java.awt.*;
 import java.util.Arrays;
 import java.util.Locale;
 import java.util.Map;
 import java.util.Set;
 import java.util.stream.Collectors;
 
-import javax.swing.JButton;
-import javax.swing.JCheckBox;
-import javax.swing.JFrame;
-import javax.swing.JPanel;
+import javax.swing.*;
 
 import cuchaz.enigma.gui.Gui;
 import cuchaz.enigma.gui.stats.StatsMember;
@@ -23,10 +20,12 @@ public class StatsDialog {
                 // init frame
                 JFrame frame = new JFrame(I18n.translate("menu.file.stats.title"));
                 JPanel checkboxesPanel = new JPanel();
+                JPanel topLevelPackagePanel = new JPanel();
                 JPanel buttonPanel = new JPanel();
-                frame.setLayout(new BorderLayout());
-                frame.add(BorderLayout.NORTH, checkboxesPanel);
-                frame.add(BorderLayout.SOUTH, buttonPanel);
+                frame.setLayout(new GridLayout(3, 0));
+                frame.add(checkboxesPanel);
+                frame.add(topLevelPackagePanel);
+                frame.add(buttonPanel);
 
                 // show checkboxes
                 Map<StatsMember, JCheckBox> checkboxes = Arrays
@@ -37,13 +36,20 @@ public class StatsDialog {
                                         return checkbox;
                                 }));
 
+                // show top-level package option
+                JLabel topLevelPackageOption = new JLabel(I18n.translate("menu.file.stats.top_level_package"));
+                JTextField topLevelPackage = new JTextField();
+                topLevelPackage.setPreferredSize(ScaleUtil.getDimension(200, 25));
+                topLevelPackagePanel.add(topLevelPackageOption);
+                topLevelPackagePanel.add(topLevelPackage);
+
                 // show generate button
                 JButton button = new JButton(I18n.translate("menu.file.stats.generate"));
                 buttonPanel.add(button);
                 button.setEnabled(false);
                 button.addActionListener(action -> {
                         frame.dispose();
-                        generateStats(gui, checkboxes);
+                        generateStats(gui, checkboxes, topLevelPackage.getText());
                 });
 
                 // add action listener to each checkbox
@@ -60,12 +66,12 @@ public class StatsDialog {
                 // show the frame
                 frame.pack();
                 frame.setVisible(true);
-                frame.setSize(ScaleUtil.getDimension(500, 120));
+                frame.setSize(ScaleUtil.getDimension(500, 150));
                 frame.setResizable(false);
                 frame.setLocationRelativeTo(gui.getFrame());
         }
 
-        private static void generateStats(Gui gui, Map<StatsMember, JCheckBox> checkboxes) {
+        private static void generateStats(Gui gui, Map<StatsMember, JCheckBox> checkboxes, String topLevelPackage) {
                 // get members from selected checkboxes
                 Set<StatsMember> includedMembers = checkboxes
                                 .entrySet()
@@ -74,9 +80,9 @@ public class StatsDialog {
                                 .map(Map.Entry::getKey)
                                 .collect(Collectors.toSet());
 
-                // checks if a projet is open
+                // checks if a project is open
                 if (gui.getController().project != null) {
-                        gui.getController().openStats(includedMembers);
+                        gui.getController().openStats(includedMembers, topLevelPackage);
                 }
         }
 }
diff --git a/enigma-swing/src/main/java/cuchaz/enigma/gui/stats/StatsGenerator.java b/enigma-swing/src/main/java/cuchaz/enigma/gui/stats/StatsGenerator.java
index aa878bb..3d031a7 100644
--- a/enigma-swing/src/main/java/cuchaz/enigma/gui/stats/StatsGenerator.java
+++ b/enigma-swing/src/main/java/cuchaz/enigma/gui/stats/StatsGenerator.java
@@ -30,7 +30,7 @@ public class StatsGenerator {
         nameProposalServices = project.getEnigma().getServices().get(NameProposalService.TYPE);
     }
 
-    public String generate(ProgressListener progress, Set<StatsMember> includedMembers) {
+    public String generate(ProgressListener progress, Set<StatsMember> includedMembers, String topLevelPackage) {
         includedMembers = EnumSet.copyOf(includedMembers);
         int totalWork = 0;
 
@@ -46,7 +46,7 @@ public class StatsGenerator {
             totalWork += entryIndex.getClasses().size();
         }
 
-        progress.init(totalWork, "progress.stats");
+        progress.init(totalWork, I18n.translate("progress.stats"));
 
         Map<String, Integer> counts = new HashMap<>();
 
@@ -97,8 +97,9 @@ public class StatsGenerator {
         Tree<Integer> tree = new Tree<>();
 
         for (Map.Entry<String, Integer> entry : counts.entrySet()) {
-            if (entry.getKey().startsWith("com.mojang")) continue; // just a few unmapped names, no point in having a subsection
-            tree.getNode(entry.getKey()).value = entry.getValue();
+            if (entry.getKey().startsWith(topLevelPackage)) {
+                tree.getNode(entry.getKey()).value = entry.getValue();
+            }
         }
 
         tree.collapse(tree.root);