Move export SRM list implementation

Author: oshando

swan-pipeline/src/main/java/de/fraunhofer/iem/swan/SwanPipeline.java

@@ -7,6 +7,7 @@
 import de.fraunhofer.iem.swan.io.dataset.SrmList;
 import de.fraunhofer.iem.swan.io.dataset.SrmListUtils;
 import de.fraunhofer.iem.swan.model.ModelEvaluator;
+import de.fraunhofer.iem.swan.model.engine.Meka;
 import de.fraunhofer.iem.swan.util.Util;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -55,12 +56,6 @@ public void run() throws IOException, InterruptedException {
         ModelEvaluator modelEvaluator = new ModelEvaluator(featuresHandler, options, testDataset.getMethods());
         modelEvaluator.trainModel();
 
-        //TODO export final list to JSON file
-        String outputFile = options.getOutputDir() + File.separator + "swan-srm-cwe-list.json";
-        ObjectMapper objectMapper = new ObjectMapper();
-        objectMapper.writeValue(new File(outputFile), dataset);
-        logger.info("SRM/CWE list exported to {}", outputFile);
-
         long analysisTime = System.currentTimeMillis() - startAnalysisTime;
         logger.info("Total runtime {} minutes", analysisTime / 60000);
     }

swan-pipeline/src/main/java/de/fraunhofer/iem/swan/data/Method.java

@@ -177,8 +177,10 @@ public Set<Category> getAllCategories() {
             return new HashSet<>();
     }
 
-    public void addCategoryClassified(Category category) {
-        this.cwe.add(category);
+    public void addCategory(Category category) {
+        if (category.isCwe())
+            cwe.add(category);
+        else srm.add(category);
     }
 
     // Inherited from SootMethodAndClass (from Soot Infoflow)

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

@@ -3,6 +3,7 @@
 import com.fasterxml.jackson.annotation.JsonIgnore;
 import de.fraunhofer.iem.swan.data.Method;
 import java.util.Set;
+import java.util.stream.Collectors;
 
 /**
  * Represents SRM JSON file.
@@ -35,4 +36,10 @@ public Set<Method> getMethods() {
     public void setMethods(Set<Method> methods) {
         this.methods = methods;
     }
+
+    public void removeUnclassifiedMethods() {
+
+        methods = methods.stream().filter(m -> m.getAllCategories().size() > 0)
+                .collect(Collectors.toSet());
+    }
 }

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

@@ -93,9 +93,10 @@ public Type appliesInternal(Method method) {
      */
     public static void exportFile(SrmList srmList, String file) throws IOException {
 
+        srmList.removeUnclassifiedMethods();
         ObjectMapper objectMapper = new ObjectMapper();
         objectMapper.writeValue(new File(file), srmList);
-        logger.info("Collected {} methods from the training set.", srmList.getMethods().size());
+        logger.info("{} methods exported to {}", srmList.getMethods().size(), file);
     }
 
     public static void removeUndocumentedMethods(SrmList list) {