Java: Put remainings parts of the generator code into the class.

Author: michaelnebel

java/ql/src/utils/model-generator/generate_flow_model.py

@@ -8,8 +8,6 @@
 import sys
 import tempfile
 
-language = "java"
-
 class Generator:
     def __init__ (self, language):
         self.language = language
@@ -18,6 +16,7 @@ def __init__ (self, language):
         self.generateSummaries = False
         self.dryRun = False
 
+
     def printHelp(self):
         print(f"""Usage:
 python3 GenerateFlowModel.py <library-database> <outputQll> [--with-sinks] [--with-sources] [--with-summaries] [--dry-run]
@@ -57,6 +56,7 @@ def setenvironment(self, target, database):
         self.workDir = tempfile.mkdtemp()
         os.makedirs(self.generatedFrameworks, exist_ok=True)
 
+
     @staticmethod
     def make(language):
         generator = Generator(language)
@@ -90,6 +90,7 @@ def make(language):
         generator.setenvironment(sys.argv[2], sys.argv[1])
         return generator
 
+
     def runQuery(self, infoMessage, query):
         print("########## Querying " + infoMessage + "...")
         queryFile = os.path.join(os.path.dirname(
@@ -147,50 +148,60 @@ def asCsvModel(self, superclass, kind, rows):
         return classTemplate.format(self.shortname[0].upper() + self.shortname[1:], kind.capitalize(), superclass, rows)
 
 
-generator = Generator.make(language)
+    def makeContent(self):
+        if self.generateSummaries:
+            summaryRows = self.runQuery("summary models", "CaptureSummaryModels.ql")
+            summaryCsv = self.asCsvModel("SummaryModelCsv", "summary", summaryRows)
+        else:
+            summaryCsv = ""
 
-if generator.generateSummaries:
-    summaryRows = generator.runQuery("summary models", "CaptureSummaryModels.ql")
-    summaryCsv = generator.asCsvModel("SummaryModelCsv", "summary", summaryRows)
-else:
-    summaryCsv = ""
+        if self.generateSinks:
+            sinkRows = self.runQuery("sink models", "CaptureSinkModels.ql")
+            sinkCsv = self.asCsvModel("SinkModelCsv", "sinks", sinkRows)
+        else:
+            sinkCsv = ""
 
-if generator.generateSinks:
-    sinkRows = generator.runQuery("sink models", "CaptureSinkModels.ql")
-    sinkCsv = generator.asCsvModel("SinkModelCsv", "sinks", sinkRows)
-else:
-    sinkCsv = ""
+        if self.generateSources:
+            sourceRows = self.runQuery("source models", "CaptureSourceModels.ql")
+            sourceCsv = self.asCsvModel("SourceModelCsv", "sources", sourceRows)
+        else:
+            sourceCsv = ""
 
-if generator.generateSources:
-    sourceRows = generator.runQuery("source models", "CaptureSourceModels.ql")
-    sourceCsv = generator.asCsvModel("SourceModelCsv", "sources", sourceRows)
-else:
-    sourceCsv = ""
+        return f"""
+/** Definitions of taint steps in the {self.shortname} framework */
 
-qllContents = f"""
-/** Definitions of taint steps in the {generator.shortname} framework */
-
-import {generator.language}
-private import semmle.code.{generator.language}.dataflow.ExternalFlow
+import {self.language}
+private import semmle.code.{self.language}.dataflow.ExternalFlow
 
 {sinkCsv}
 {sourceCsv}
 {summaryCsv}
 
-"""
+        """
 
-if generator.dryRun:
-    print("CSV Models generated, but not written to file.")
-    sys.exit(0)
 
-with open(generator.frameworkTarget, "w") as frameworkQll:
-    frameworkQll.write(qllContents)
+    def save(self, content):
+        with open(self.frameworkTarget, "w") as frameworkQll:
+            frameworkQll.write(content)
 
-cmd = ['codeql', 'query', 'format', '--in-place', generator.frameworkTarget]
-ret = subprocess.call(cmd)
-if ret != 0:
-    print("Failed to format query. Failed command was: " + shlex.join(cmd))
-    sys.exit(1)
+        cmd = ['codeql', 'query', 'format', '--in-place', self.frameworkTarget]
+        ret = subprocess.call(cmd)
+        if ret != 0:
+            print("Failed to format query. Failed command was: " + shlex.join(cmd))
+            sys.exit(1)
+
+        print("")
+        print("CSV model written to " + self.frameworkTarget)
+
+
+    def run(self):
+        content = self.makeContent()
 
-print("")
-print("CSV model written to " + generator.frameworkTarget)
+        if self.dryRun:
+            print("CSV Models generated, but not written to file.")
+            sys.exit(0)
+        
+        self.save(content)
+
+language = "java"
+Generator.make(language).run()