Skip to content

Commit 95c94ef

Browse files
committed
MaD: update generator scripts to support writing JSON data extensions
The `generate_mad.py` script accepts a flag that lets you choose between YAML or JSON, but YAML is still the default. We can change the default later, when CLI support for JSON is widely available. `convert_extensions.py` also still hard-codes `.yml` as the output extension/format.
1 parent e8e8657 commit 95c94ef

4 files changed

Lines changed: 48 additions & 8 deletions

File tree

‎misc/scripts/models-as-data/bulk_generate_mad.py‎

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -430,10 +430,11 @@ def download_and_decompress(analyzed_database: dict) -> str:
430430
def clean_up_mad_destination_for_project(config, name: str):
431431
target = pathlib.Path(config["destination"], name)
432432
if config.get("single-file", False):
433-
target = target.with_suffix(".model.yml")
434-
if target.exists():
435-
print(f"Deleting existing MaD file at {target}")
436-
target.unlink()
433+
for extension in (".model.json", ".model.yml"):
434+
model_target = target.with_suffix(extension)
435+
if model_target.exists():
436+
print(f"Deleting existing MaD file at {model_target}")
437+
model_target.unlink()
437438
elif target.exists():
438439
print(f"Deleting existing MaD directory at {target}")
439440
shutil.rmtree(target, ignore_errors=True)

‎misc/scripts/models-as-data/data_extension_file.py‎

Lines changed: 29 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
import json
22

33

4-
# Simple class for representing a data extension file and writing it to disk in a pretty-printed YAML format.
4+
# Simple class for representing a data extension file and writing it to disk in a pretty-printed YAML or JSON format.
55
class DataExtensionFile:
66
def __init__(self, pack):
77
self.pack = pack
@@ -38,3 +38,31 @@ def write_yaml(self, f):
3838
f.write("extensions:\n")
3939
for predicate in sorted(self.predicates.keys()):
4040
f.write(self.yaml_for_predicate(predicate))
41+
42+
def json_for_row(self, row):
43+
return "[" + ", ".join(json.dumps(value) for value in row) + "]"
44+
45+
def json_for_predicate(self, predicate):
46+
return f""" {{
47+
"addsTo": {{
48+
"pack": "{self.pack}",
49+
"extensible": "{predicate}"
50+
}},
51+
"data": [
52+
{',\n '.join(self.json_for_row(row) for row in self.predicates[predicate])}
53+
]
54+
}}"""
55+
56+
def write_json(self, f):
57+
# It would be preferable to use a serializer like Python's json.dumps, but it can't
58+
# pretty-print with one line per tuple.
59+
f.write("// THIS FILE IS AN AUTO-GENERATED MODELS AS DATA FILE. DO NOT EDIT.\n")
60+
f.write("{\n")
61+
f.write(' "extensions": [\n')
62+
for i, predicate in enumerate(self.predicates):
63+
if i > 0:
64+
f.write(",\n")
65+
f.write(self.json_for_predicate(predicate))
66+
f.write("\n")
67+
f.write(" ]\n")
68+
f.write("}")

‎misc/scripts/models-as-data/generate_mad.py‎

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -119,6 +119,12 @@ def make():
119119
"--single-file",
120120
help="Generate a single file with all models instead of separate files for each namespace, using provided argument as the base filename.",
121121
)
122+
p.add_argument(
123+
"--extension-format",
124+
choices=["json", "yaml"],
125+
default="yaml",
126+
help="Format for the generated data extension files (default %(default)s)",
127+
)
122128
generator = p.parse_args(namespace=Generator())
123129

124130
if (
@@ -222,17 +228,19 @@ def run(self):
222228
print("Models as data extensions generated, but not written to file.")
223229
sys.exit(0)
224230

231+
file_extension = ".json" if self.extension_format == "json" else ".yml"
232+
225233
if (
226234
self.with_sinks
227235
or self.with_sources
228236
or self.with_summaries
229237
or self.with_neutrals
230238
):
231-
merger.save(self.generated_frameworks, ".model.yml")
239+
merger.save(self.generated_frameworks, f".model{file_extension}")
232240

233241
if self.with_typebased_summaries:
234242
type_based_extensions.save(
235-
self.generated_frameworks, ".typebased.model.yml"
243+
self.generated_frameworks, f".typebased.model{file_extension}"
236244
)
237245

238246

‎misc/scripts/models-as-data/helpers.py‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -51,7 +51,10 @@ def save(self, dir, file_extension):
5151
)
5252
target = os.path.join(dir, f"{sanitized_namespace}{file_extension}")
5353
with open(target, "w") as f:
54-
extension.write_yaml(f)
54+
if file_extension.endswith(".json"):
55+
extension.write_json(f)
56+
else:
57+
extension.write_yaml(f)
5558
print("Models as data extensions written to " + target)
5659

5760

0 commit comments

Comments
 (0)