Bug 1886702 - Fix util.load_yaml_or_json failures when loading json files

glean_parser/util.py

@@ -10,7 +10,7 @@
 from pathlib import Path
 import sys
 import textwrap
-from typing import Any, Callable, Iterable, Sequence, Tuple, Union, Optional
+from typing import Any, Callable, Iterable, Sequence, Tuple, Union, Optional, TextIO
 import urllib.request
 
 import appdirs  # type: ignore
@@ -97,6 +97,18 @@ def _construct_mapping_adding_line(loader, node):
     return yaml.load(stream, SafeLineLoader)
 
 
+def json_load(stream: TextIO):
+
+    # The Python json parser doesn't give us access to the line number
+    def _construct_mapping_adding_empty_line(obj):
+        mapping = DictWrapper(obj)
+        mapping.defined_in = {"line": 0}
+        return mapping
+
+    return json.load(stream, object_hook=_construct_mapping_adding_empty_line)
+
+
+
 def ordered_yaml_dump(data, **kwargs):
     class OrderedDumper(yaml.Dumper):
         pass
@@ -126,7 +138,7 @@ def load_yaml_or_json(path: Path):
 
     if path.suffix == ".json":
         with path.open("r", encoding="utf-8") as fd:
-            return json.load(fd)
+            return json_load(fd)
     elif path.suffix in (".yml", ".yaml", ".yamlx"):
         with path.open("r", encoding="utf-8") as fd:
             return yaml_load(fd)

tests/test_parser.py

@@ -5,13 +5,16 @@
 
 from pathlib import Path
 import json
+import os
 import re
 import textwrap
+import tempfile
 
 import pytest
 
 from glean_parser import metrics
 from glean_parser import parser
+from glean_parser.util import load_yaml_or_json
 
 import util
 
@@ -1181,6 +1184,21 @@ def test_no_internal_fields_exposed():
     assert expected_json == out_json
 
 
+def test_json_loader():
+    yaml_input = ROOT / "data" / "all_metrics.yaml"
+    yaml_obj = load_yaml_or_json(yaml_input)
+
+    with tempfile.NamedTemporaryFile(suffix=".json", mode="w+", delete=False) as fd:
+        json.dump(yaml_obj, fd)
+        json_input = Path(fd.name)
+
+    json_obj = load_yaml_or_json(json_input)
+
+    os.remove(json_input)
+
+    assert yaml_obj == json_obj
+
+
 def test_object():
     structure = {"type": "array", "items": {"type": "number"}}
     contents = [{"category": {"metric": {"type": "object", "structure": structure}}}]