99 lines
3.4 KiB
Python
99 lines
3.4 KiB
Python
# Copyright (c) 2022 Tulir Asokan
|
|
#
|
|
# This Source Code Form is subject to the terms of the Mozilla Public
|
|
# License, v. 2.0. If a copy of the MPL was not distributed with this
|
|
# file, You can obtain one at http://mozilla.org/MPL/2.0/.
|
|
from __future__ import annotations
|
|
|
|
from typing import Any, Generic, Type, TypeVar
|
|
import copy
|
|
|
|
from ruamel.yaml.comments import CommentedMap
|
|
|
|
T = TypeVar("T")
|
|
|
|
|
|
class RecursiveDict(Generic[T]):
|
|
def __init__(self, data: T | None = None, dict_factory: Type[T] | None = None) -> None:
|
|
self._dict_factory = dict_factory or dict
|
|
self._data: CommentedMap = data or self._dict_factory()
|
|
|
|
def clone(self) -> RecursiveDict:
|
|
return RecursiveDict(data=copy.deepcopy(self._data), dict_factory=self._dict_factory)
|
|
|
|
@staticmethod
|
|
def parse_key(key: str) -> tuple[str, str | None]:
|
|
if "." not in key:
|
|
return key, None
|
|
key, next_key = key.split(".", 1)
|
|
if len(key) > 0 and key[0] == "[":
|
|
end_index = next_key.index("]")
|
|
key = key[1:] + "." + next_key[:end_index]
|
|
next_key = next_key[end_index + 2 :] if len(next_key) > end_index + 1 else None
|
|
return key, next_key
|
|
|
|
def _recursive_get(self, data: T, key: str, default_value: Any) -> Any:
|
|
key, next_key = self.parse_key(key)
|
|
if next_key is not None:
|
|
next_data = data.get(key, self._dict_factory())
|
|
return self._recursive_get(next_data, next_key, default_value)
|
|
try:
|
|
return data[key]
|
|
except (AttributeError, KeyError):
|
|
return default_value
|
|
|
|
def get(self, key: str, default_value: Any, allow_recursion: bool = True) -> Any:
|
|
if allow_recursion and "." in key:
|
|
return self._recursive_get(self._data, key, default_value)
|
|
return self._data.get(key, default_value)
|
|
|
|
def __getitem__(self, key: str) -> Any:
|
|
return self.get(key, None)
|
|
|
|
def __contains__(self, key: str) -> bool:
|
|
return self.get(key, None) is not None
|
|
|
|
def _recursive_set(self, data: T, key: str, value: Any) -> None:
|
|
key, next_key = self.parse_key(key)
|
|
if next_key is not None:
|
|
if key not in data:
|
|
data[key] = self._dict_factory()
|
|
next_data = data.get(key, self._dict_factory())
|
|
return self._recursive_set(next_data, next_key, value)
|
|
data[key] = value
|
|
|
|
def set(self, key: str, value: Any, allow_recursion: bool = True) -> None:
|
|
if allow_recursion and "." in key:
|
|
self._recursive_set(self._data, key, value)
|
|
return
|
|
self._data[key] = value
|
|
|
|
def __setitem__(self, key: str, value: Any) -> None:
|
|
self.set(key, value)
|
|
|
|
def _recursive_del(self, data: T, key: str) -> None:
|
|
key, next_key = self.parse_key(key)
|
|
if next_key is not None:
|
|
if key not in data:
|
|
return
|
|
next_data = data[key]
|
|
return self._recursive_del(next_data, next_key)
|
|
try:
|
|
del data[key]
|
|
del data.ca.items[key]
|
|
except KeyError:
|
|
pass
|
|
|
|
def delete(self, key: str, allow_recursion: bool = True) -> None:
|
|
if allow_recursion and "." in key:
|
|
self._recursive_del(self._data, key)
|
|
return
|
|
try:
|
|
del self._data[key]
|
|
del self._data.ca.items[key]
|
|
except KeyError:
|
|
pass
|
|
|
|
def __delitem__(self, key: str) -> None:
|
|
self.delete(key)
|