@@ -8,25 +8,6 @@
 import tempfile
 
 
-def quote_if_needed(v):
-    # string columns
-    if type(v) is str:
-        return '"' + v + '"'
-    # bool column
-    return str(v)
-
-
-def parseData(data):
-    rows = [{}, {}]
-    for row in data:
-        d = map(quote_if_needed, row)
-        provenance = row[-1]
-        targetRows = rows[1] if provenance.endswith("generated") else rows[0]
-        helpers.insert_update(targetRows, row[0], "      - [" + ", ".join(d) + "]\n")
-
-    return rows
-
-
 class Converter:
     def __init__(self, language, dbDir):
         self.language = language
@@ -63,48 +44,38 @@ def runQuery(self, query):
         )
         return helpers.readData(self.workDir, resultBqrs)
 
-    def asAddsTo(self, rows, predicate):
-        extensions = [{}, {}]
-        for i in range(2):
-            for key in rows[i]:
-                extensions[i][key] = helpers.addsToTemplate.format(
-                    f"codeql/{self.language}-all", predicate, rows[i][key]
-                )
-
-        return extensions
-
-    def getAddsTo(self, query, predicate):
+    def merge_query_results(self, query, predicate, mergers):
         data = self.runQuery(query)
-        rows = parseData(data)
-        return self.asAddsTo(rows, predicate)
-
-    def makeContent(self):
-        summaries = self.getAddsTo("ExtractSummaries.ql", helpers.summaryModelPredicate)
-        sources = self.getAddsTo("ExtractSources.ql", helpers.sourceModelPredicate)
-        sinks = self.getAddsTo("ExtractSinks.ql", helpers.sinkModelPredicate)
-        neutrals = self.getAddsTo("ExtractNeutrals.ql", helpers.neutralModelPredicate)
-        return [
-            helpers.merge(sources[0], sinks[0], summaries[0], neutrals[0]),
-            helpers.merge(sources[1], sinks[1], summaries[1], neutrals[1]),
+        for row in data:
+            provenance = row[-1]
+            namespace = row[0]
+            target_merger = (
+                mergers[1] if provenance.endswith("generated") else mergers[0]
+            )
+            target_merger.add_row(namespace, predicate, row)
+
+    def make_extensions(self):
+        mergers = [
+            helpers.ExtensionMerger(f"codeql/{self.language}-all"),
+            helpers.ExtensionMerger(f"codeql/{self.language}-all"),
         ]
+        self.merge_query_results(
+            "ExtractSummaries.ql", helpers.summaryModelPredicate, mergers
+        )
+        self.merge_query_results(
+            "ExtractSources.ql", helpers.sourceModelPredicate, mergers
+        )
+        self.merge_query_results("ExtractSinks.ql", helpers.sinkModelPredicate, mergers)
+        self.merge_query_results(
+            "ExtractNeutrals.ql", helpers.neutralModelPredicate, mergers
+        )
+        return mergers
+
+    def run(self):
+        mergers = self.make_extensions()
 
-    def save(self, extensions):
         # Create directory if it doesn't exist
         os.makedirs(self.extDir, exist_ok=True)
 
-        # Create a file for each namespace and save models.
-        extensionTemplate = """extensions:
-{0}"""
-        for entry in extensions[0]:
-            with open(self.extDir + "/" + entry + self.modelFileExtension, "w") as f:
-                f.write(extensionTemplate.format(extensions[0][entry]))
-
-        for entry in extensions[1]:
-            with open(
-                self.extDir + "/generated/" + entry + self.modelFileExtension, "w"
-            ) as f:
-                f.write(extensionTemplate.format(extensions[1][entry]))
-
-    def run(self):
-        extensions = self.makeContent()
-        self.save(extensions)
+        mergers[0].save(self.extDir, self.modelFileExtension)
+        mergers[1].save(self.extDir + "/generated", self.modelFileExtension)
@@ -0,0 +1,48 @@
+import json
+
+
+class DataExtensionFile:
+    """Simple class for representing a data extension file and writing it to disk in a
+    pretty-printed YAML format."""
+
+    def __init__(self, pack):
+        self.pack = pack
+        # predicate -> dict of tuples (used as an ordered set)
+        self.predicates = {}
+
+    def add_rows(self, predicate, rows):
+        """Add the rows, which can be any collection of any collection of string/boolean values."""
+        if predicate not in self.predicates:
+            self.predicates[predicate] = dict()
+        d = self.predicates[predicate]
+        for row in rows:
+            # To be agnostic about the type of collection used for rows, convert each one to a
+            # tuple. Insert the tuple as a key in the dict, associating it with a dummy value.
+            d[tuple(row)] = None
+
+    def yaml_for_value(self, value):
+        # if is boolean:
+        if isinstance(value, bool):
+            # json.dumps would produce lower-case "true" or "false". Even though that's valid YAML,
+            # we have historically used the upper-case variants.
+            return "True" if value else "False"
+        return json.dumps(value)
+
+    def yaml_for_row(self, row):
+        return f"[{', '.join(self.yaml_for_value(value) for value in row)}]"
+
+    def yaml_for_predicate(self, predicate):
+        return f"""  - addsTo:
+      pack: {self.pack}
+      extensible: {predicate}
+    data:
+      - {'\n      - '.join(self.yaml_for_row(row) for row in self.predicates[predicate].keys())}
+"""
+
+    def write_yaml(self, f):
+        # It would be preferable to use a YAML serialization library, but this lets us control the
+        # formatting e.g. to produce one line per tuple.
+        f.write("# THIS FILE IS AN AUTO-GENERATED MODELS AS DATA FILE. DO NOT EDIT.\n")
+        f.write("extensions:\n")
+        for predicate in sorted(self.predicates.keys()):
+            f.write(self.yaml_for_predicate(predicate))
@@ -6,29 +6,8 @@
 import subprocess
 import sys
 import tempfile
-import re
 import argparse
 
-
-def quote_if_needed(row):
-    if row != "true" and row != "false":
-        return '"' + row + '"'
-    # subtypes column
-    return row[0].upper() + row[1:]
-
-
-def parseData(data):
-    rows = {}
-
-    for row in data:
-        d = row[0].split(";")
-        namespace = d[0]
-        d = map(quote_if_needed, d)
-        helpers.insert_update(rows, namespace, "      - [" + ", ".join(d) + "]\n")
-
-    return rows
-
-
 description = """\
 This generates summary, source, sink and neutral models for the code in the database.
 The files will be placed in `LANGUAGE/ql/lib/ext/generated/modelgenerator/DIR`"""
@@ -182,76 +161,65 @@ def runQuery(self, query):
 
         return helpers.readData(self.workDir, resultBqrs)
 
-    def asAddsTo(self, rows, predicate):
-        extensions = {}
-        for key in rows:
-            extensions[key] = helpers.addsToTemplate.format(
-                f"codeql/{self.language}-all", predicate, rows[key]
-            )
-        return extensions
-
-    def getAddsTo(self, query, predicate):
+    def convert_if_needed(self, value):
+        # Convert a string value to a boolean if it is "true" or "false". Otherwise, return the original
+        # value as a string. Since no extensible predicate currently uses integers, don't attempt to
+        # convert numeric strings to integers.
+        if value == "true":
+            return True
+        if value == "false":
+            return False
+        return value
+
+    def merge_query_results(self, query, predicate, merger):
         data = self.runQuery(query)
-        rows = parseData(data)
-        if self.single_file and rows:
-            rows = {self.single_file: "".join(rows.values())}
-        return self.asAddsTo(rows, predicate)
+        for row in data:
+            d = row[0].split(";")
+            namespace = d[0]
+            d = list(map(self.convert_if_needed, d))
+            if self.single_file:
+                merger.add_row(self.single_file, predicate, d)
+            else:
+                merger.add_row(namespace, predicate, d)
+
+    def make_extensions(self):
+        merger = helpers.ExtensionMerger(f"codeql/{self.language}-all")
 
-    def makeContent(self):
-        summaryAddsTo = {}
         if self.with_summaries:
-            summaryAddsTo = self.getAddsTo(
-                "CaptureSummaryModels.ql", helpers.summaryModelPredicate
+            self.merge_query_results(
+                "CaptureSummaryModels.ql", helpers.summaryModelPredicate, merger
             )
 
-        sinkAddsTo = {}
         if self.with_sinks:
-            sinkAddsTo = self.getAddsTo(
-                "CaptureSinkModels.ql", helpers.sinkModelPredicate
+            self.merge_query_results(
+                "CaptureSinkModels.ql", helpers.sinkModelPredicate, merger
             )
 
-        sourceAddsTo = {}
         if self.with_sources:
-            sourceAddsTo = self.getAddsTo(
-                "CaptureSourceModels.ql", helpers.sourceModelPredicate
+            self.merge_query_results(
+                "CaptureSourceModels.ql", helpers.sourceModelPredicate, merger
             )
 
-        neutralAddsTo = {}
         if self.with_neutrals:
-            neutralAddsTo = self.getAddsTo(
-                "CaptureNeutralModels.ql", helpers.neutralModelPredicate
+            self.merge_query_results(
+                "CaptureNeutralModels.ql", helpers.neutralModelPredicate, merger
             )
 
-        return helpers.merge(summaryAddsTo, sinkAddsTo, sourceAddsTo, neutralAddsTo)
+        return merger
 
-    def makeTypeBasedContent(self):
+    def make_type_based_extensions(self):
+        merger = helpers.ExtensionMerger(f"codeql/{self.language}-all")
         if self.with_typebased_summaries:
-            typeBasedSummaryAddsTo = self.getAddsTo(
-                "CaptureTypeBasedSummaryModels.ql", helpers.summaryModelPredicate
-            )
-        else:
-            typeBasedSummaryAddsTo = {}
-
-        return typeBasedSummaryAddsTo
-
-    def save(self, extensions, extension):
-        # Create a file for each namespace and save models.
-        extensionTemplate = """# THIS FILE IS AN AUTO-GENERATED MODELS AS DATA FILE. DO NOT EDIT.
-extensions:
-{0}"""
-        for entry in extensions:
-            # Replace problematic characters with dashes, and collapse multiple dashes.
-            sanitizedEntry = re.sub(
-                r"-+", "-", entry.replace("/", "-").replace(":", "-")
+            self.merge_query_results(
+                "CaptureTypeBasedSummaryModels.ql",
+                helpers.summaryModelPredicate,
+                merger,
             )
-            target = os.path.join(self.generated_frameworks, sanitizedEntry + extension)
-            with open(target, "w") as f:
-                f.write(extensionTemplate.format(extensions[entry]))
-            print("Models as data extensions written to " + target)
+        return merger
 
     def run(self):
-        content = self.makeContent()
-        typeBasedContent = self.makeTypeBasedContent()
+        merger = self.make_extensions()
+        type_based_extensions = self.make_type_based_extensions()
 
         if self.dry_run:
             print("Models as data extensions generated, but not written to file.")
@@ -263,10 +231,12 @@ def run(self):
             or self.with_summaries
             or self.with_neutrals
         ):
-            self.save(content, ".model.yml")
+            merger.save(self.generated_frameworks, ".model.yml")
 
         if self.with_typebased_summaries:
-            self.save(typeBasedContent, ".typebased.model.yml")
+            type_based_extensions.save(
+                self.generated_frameworks, ".typebased.model.yml"
+            )
 
 
 if __name__ == "__main__":
 
@@ -2,17 +2,53 @@
 import os
 import shutil
 import subprocess
+import re
+import data_extension_file
 
 # Shared strings.
 summaryModelPredicate = "summaryModel"
 sinkModelPredicate = "sinkModel"
 sourceModelPredicate = "sourceModel"
 neutralModelPredicate = "neutralModel"
-addsToTemplate = """  - addsTo:
-      pack: {0}
-      extensible: {1}
-    data:
-{2}"""
+
+
+class ExtensionMerger:
+    """Helper class for accumulating tuples grouped by namespace and predicate, and generating data
+    extensions for them."""
+
+    def __init__(self, pack):
+        self.pack = pack
+        # namespace -> predicate -> dict of tuples (used as an ordered set)
+        self.namespaces = {}
+
+    def add_row(self, namespace, predicate, row):
+        """Add the row, which can be any type of collection of string/boolean values, to the set of
+        tuples for the given namespace and predicate."""
+        if namespace not in self.namespaces:
+            self.namespaces[namespace] = {}
+        if predicate not in self.namespaces[namespace]:
+            self.namespaces[namespace][predicate] = dict()
+        self.namespaces[namespace][predicate][tuple(row)] = None
+
+    def save(self, dir, file_extension):
+        """Save the accumulated rows as data extension files in the specified directory. Each
+        namespace will result in a separate file, with the file name derived from the namespace,
+        plus the given file extension."""
+        # Create a file for each namespace and save models.
+        for namespace in self.namespaces:
+            extension = data_extension_file.DataExtensionFile(self.pack)
+            for predicate in self.namespaces[namespace]:
+                extension.add_rows(
+                    predicate, self.namespaces[namespace][predicate].keys()
+                )
+            # Replace problematic characters with dashes, and collapse multiple dashes.
+            sanitized_namespace = re.sub(
+                r"-+", "-", namespace.replace("/", "-").replace(":", "-")
+            )
+            target = os.path.join(dir, f"{sanitized_namespace}{file_extension}")
+            with open(target, "w") as f:
+                extension.write_yaml(f)
+            print("Models as data extensions written to " + target)
 
 
 def remove_dir(dirName):