Move dataset filters to data processor

Author: oshando

swan-pipeline/src/main/java/de/fraunhofer/iem/swan/io/dataset/DatasetProcessor.java

@@ -1,13 +1,15 @@
 package de.fraunhofer.iem.swan.io.dataset;
 
 import de.fraunhofer.iem.swan.cli.SwanOptions;
+import de.fraunhofer.iem.swan.data.Method;
 import de.fraunhofer.iem.swan.io.doc.JavadocProcessor;
 import de.fraunhofer.iem.swan.soot.Soot;
 import de.fraunhofer.iem.swan.util.Util;
+import edu.stanford.nlp.util.StringUtils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
-
 import java.io.IOException;
+import java.util.HashSet;
 
 public class DatasetProcessor {
 
@@ -31,23 +33,44 @@ public Dataset run() {
 
             if (!options.getTrainDataDir().isEmpty())
                 soot.cleanupList(dataset.getTrain());
+
+            logger.info("Importing {} SRMs from TRAIN dataset in {}, distribution={}",
+                    dataset.getTrainMethods().size(), options.getDatasetJson(),
+                    Util.countCategories(dataset.getTrainMethods()));
+
+            //Apply filters to dataset
+            if (options.getDiscovery().size() > 0 || options.isDocumented()) {
+
+                for (Method method : new HashSet<>(dataset.getTrainMethods())) {
+
+                    if (!options.getDiscovery().contains(method.getDiscovery()) ||
+                            ((method.getJavadoc().getMethodComment().length() == 0
+                                    || StringUtils.split(method.getJavadoc().getMethodComment(), " ").size() <= 1) && options.isDocumented())) {
+                        dataset.getTrainMethods().remove(method);
+                    }
+                }
+            }
+
         } catch (IOException e) {
             throw new RuntimeException(e);
         }
 
-        logger.info("Loaded {} training methods, distribution={}", dataset.getTrainMethods().size(), Util.countCategories(dataset.getTrainMethods()));
-
         if (options.getPhase().equals("predict")) {
             //Load methods from the test set
             dataset.setTest(new SrmList(options.getTestDataDir()));
             dataset.getTest().setMethods(soot.loadMethods(dataset.getTest().getTestClasses()));
 
+            logger.info("Importing {} SRMs from TEST dataset in {}, distribution={}",
+                    dataset.getTestMethods().size(), options.getTestDataDir(),
+                    Util.countCategories(dataset.getTestMethods()));
+
             if (options.getFeatureSet().contains("doc-")) {
 
                 //Extract doc comments and add to test set, if option is selected
                 JavadocProcessor javadocProcessor = new JavadocProcessor(options.getTestDataSourceDir(), options.getOutputDir());
                 javadocProcessor.run(dataset.getTestMethods(), options.getFeatureSet());
-                logger.info("Loaded {} methods from {}", dataset.getTestMethods().size(), options.getTestDataDir());
+
+                logger.info("Extracting doc comments for {} methods in {}", dataset.getTestMethods().size(), options.getTestDataDir());
             }
         }
         return dataset;

swan-pipeline/src/main/java/de/fraunhofer/iem/swan/io/dataset/SrmListUtils.java

@@ -4,15 +4,12 @@
 import de.fraunhofer.iem.swan.data.Method;
 import de.fraunhofer.iem.swan.io.doc.Javadoc;
 import de.fraunhofer.iem.swan.io.doc.ssldoclet.MethodBlockType;
-import edu.stanford.nlp.util.StringUtils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 import java.io.File;
 import java.io.IOException;
 import java.util.ArrayList;
-import java.util.HashSet;
-import java.util.List;
 import java.util.Set;
 
 /**
@@ -34,11 +31,8 @@ public class SrmListUtils {
     public static SrmList importFile(String file) throws IOException {
 
         ObjectMapper objectMapper = new ObjectMapper();
-        SrmList srmList = objectMapper.readValue(new File(file), SrmList.class);
-        logger.info("Collected {} methods from the training set.", srmList.getMethods().size());
 
-      //  removeManualMethods(srmList);
-        return srmList;
+        return objectMapper.readValue(new File(file), SrmList.class);
     }
 
     /**
@@ -51,38 +45,18 @@ public static void exportFile(SrmList srmList, String file) throws IOException {
 
         ObjectMapper objectMapper = new ObjectMapper();
         objectMapper.writeValue(new File(file), srmList);
-        logger.info("{} SRMs exported to {}", srmList.getMethods().size(), file);
-    }
-
-    public static void removeUndocumentedMethods(SrmList list) {
-        Set<Method> temp = new HashSet<>(list.getMethods());
-
-        for (Method method : temp) {
-            List<String> words = StringUtils.split(method.getJavadoc().getMethodComment(), " ");
-
-            if (method.getJavadoc().getMethodComment().length() == 0 || words.size() <= 1)
-                list.getMethods().remove(method);
-        }
-    }
-
-    public static void removeManualMethods(SrmList list) {
-        Set<Method> temp = new HashSet<>(list.getMethods());
-
-        for (Method method : temp) {
-
-            if(!method.getDiscovery().contains("manual"))
-                list.getMethods().remove(method);
-        }
+        logger.info("Exporting {} SRMs to {}", srmList.getMethods().size(), file);
     }
 
     /**
      * Adds doc comments to method set.
-     * @param methods methods to be updated
+     *
+     * @param methods  methods to be updated
      * @param javadocs list of Javadoc objects
      * @return updated method set
      */
     public static Set<Method> addDocComments(Set<Method> methods, ArrayList<Javadoc> javadocs) {
-
+        //TODO Check if returned set is used
         for (Javadoc doc : javadocs) {
 
             for (MethodBlockType methodBlock : doc.getMethodBlocks().values()) {