Skip to content

Commit

Permalink
Revert "Rewrite yaml loader (#13803)"
Browse files Browse the repository at this point in the history
This reverts commit 38ff46e.
  • Loading branch information
hawkeye217 authored Sep 17, 2024
1 parent 1ed8642 commit 8994edd
Show file tree
Hide file tree
Showing 4 changed files with 29 additions and 30 deletions.
2 changes: 1 addition & 1 deletion .gitignore
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
.DS_Store
__pycache__
*.pyc
*.swp
debug
.vscode/*
Expand Down
7 changes: 3 additions & 4 deletions frigate/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,6 @@
from typing import Any, Dict, List, Optional, Tuple, Union

import numpy as np
import yaml
from pydantic import (
BaseModel,
ConfigDict,
Expand Down Expand Up @@ -42,11 +41,11 @@
)
from frigate.plus import PlusApi
from frigate.util.builtin import (
NoDuplicateKeysLoader,
deep_merge,
escape_special_characters,
generate_color_palette,
get_ffmpeg_arg_list,
load_config_with_no_duplicates,
)
from frigate.util.config import StreamInfoRetriever, get_relative_coordinates
from frigate.util.image import create_mask
Expand Down Expand Up @@ -1765,13 +1764,13 @@ def parse_file(cls, config_file):
raw_config = f.read()

if config_file.endswith(YAML_EXT):
config = yaml.load(raw_config, NoDuplicateKeysLoader)
config = load_config_with_no_duplicates(raw_config)
elif config_file.endswith(".json"):
config = json.loads(raw_config)

return cls.model_validate(config)

@classmethod
def parse_raw(cls, raw_config):
config = yaml.load(raw_config, NoDuplicateKeysLoader)
config = load_config_with_no_duplicates(raw_config)
return cls.model_validate(config)
5 changes: 2 additions & 3 deletions frigate/test/test_config.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,14 +4,13 @@
from unittest.mock import patch

import numpy as np
import yaml
from pydantic import ValidationError

from frigate.config import BirdseyeModeEnum, FrigateConfig
from frigate.const import MODEL_CACHE_DIR
from frigate.detectors import DetectorTypeEnum
from frigate.plus import PlusApi
from frigate.util.builtin import NoDuplicateKeysLoader, deep_merge
from frigate.util.builtin import deep_merge, load_config_with_no_duplicates


class TestConfig(unittest.TestCase):
Expand Down Expand Up @@ -1538,7 +1537,7 @@ def test_fails_duplicate_keys(self):
"""

self.assertRaises(
ValueError, lambda: yaml.load(raw_config, NoDuplicateKeysLoader)
ValueError, lambda: load_config_with_no_duplicates(raw_config)
)

def test_object_filter_ratios_work(self):
Expand Down
45 changes: 23 additions & 22 deletions frigate/util/builtin.py
Original file line number Diff line number Diff line change
Expand Up @@ -89,31 +89,32 @@ def deep_merge(dct1: dict, dct2: dict, override=False, merge_lists=False) -> dic
return merged


class NoDuplicateKeysLoader(yaml.loader.SafeLoader):
"""A yaml SafeLoader that disallows duplicate keys"""

def construct_mapping(self, node, deep=False):
mapping = super().construct_mapping(node, deep=deep)

if len(node.value) != len(mapping):
# There's a duplicate key somewhere. Find it.
duplicate_keys = [
key
for key, count in Counter(
self.construct_object(key, deep=deep) for key, _ in node.value
)
if count > 1
]
def load_config_with_no_duplicates(raw_config) -> dict:
"""Get config ensuring duplicate keys are not allowed."""

# This might be possible if PyYAML's construct_mapping() changes the node
# afterwards for some reason? I don't see why, but better safe than sorry.
assert len(duplicate_keys) > 0
# https://stackoverflow.com/a/71751051
# important to use SafeLoader here to avoid RCE
class PreserveDuplicatesLoader(yaml.loader.SafeLoader):
pass

raise ValueError(
f"Config field duplicates are not allowed, the following fields are duplicated in the config: {', '.join(duplicate_keys)}"
)
def map_constructor(loader, node, deep=False):
keys = [loader.construct_object(node, deep=deep) for node, _ in node.value]
vals = [loader.construct_object(node, deep=deep) for _, node in node.value]
key_count = Counter(keys)
data = {}
for key, val in zip(keys, vals):
if key_count[key] > 1:
raise ValueError(
f"Config input {key} is defined multiple times for the same field, this is not allowed."
)
else:
data[key] = val
return data

return mapping
PreserveDuplicatesLoader.add_constructor(
yaml.resolver.BaseResolver.DEFAULT_MAPPING_TAG, map_constructor
)
return yaml.load(raw_config, PreserveDuplicatesLoader)


def clean_camera_user_pass(line: str) -> str:
Expand Down

0 comments on commit 8994edd

Please sign in to comment.