Skip to content

Commit

Permalink
Rewrite yaml loader (#13803)
Browse files Browse the repository at this point in the history
* Ignore entire __pycache__ folder instead of individual *.pyc files

* Rewrite the yaml loader to match PyYAML

The old implementation would fail in weird ways with configs that were
incorrect in just the right way. The new implementation just does what
PyYAML would do, only diverging in case of duplicate keys.

* Clarify duplicate yaml key ValueError message

Co-authored-by: Nicolas Mowen <nickmowen213@gmail.com>

---------

Co-authored-by: Nicolas Mowen <nickmowen213@gmail.com>
  • Loading branch information
gtsiam and NickM-27 committed Sep 17, 2024
1 parent 2362d0e commit 38ff46e
Show file tree
Hide file tree
Showing 4 changed files with 30 additions and 29 deletions.
2 changes: 1 addition & 1 deletion .gitignore
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
.DS_Store
*.pyc
__pycache__
*.swp
debug
.vscode/*
Expand Down
7 changes: 4 additions & 3 deletions frigate/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@
from typing import Any, Dict, List, Optional, Tuple, Union

import numpy as np
import yaml
from pydantic import (
BaseModel,
ConfigDict,
Expand Down Expand Up @@ -41,11 +42,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 @@ -1764,13 +1765,13 @@ def parse_file(cls, config_file):
raw_config = f.read()

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

return cls.model_validate(config)

@classmethod
def parse_raw(cls, raw_config):
config = load_config_with_no_duplicates(raw_config)
config = yaml.load(raw_config, NoDuplicateKeysLoader)
return cls.model_validate(config)
5 changes: 3 additions & 2 deletions frigate/test/test_config.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,13 +4,14 @@
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 deep_merge, load_config_with_no_duplicates
from frigate.util.builtin import NoDuplicateKeysLoader, deep_merge


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

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

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


def load_config_with_no_duplicates(raw_config) -> dict:
"""Get config ensuring duplicate keys are not allowed."""
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
]

# https://stackoverflow.com/a/71751051
# important to use SafeLoader here to avoid RCE
class PreserveDuplicatesLoader(yaml.loader.SafeLoader):
pass
# 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

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
raise ValueError(
f"Config field duplicates are not allowed, the following fields are duplicated in the config: {', '.join(duplicate_keys)}"
)

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


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

0 comments on commit 38ff46e

Please sign in to comment.