from typing import Any, Dict, List, Optional, Type
from qdrant_client.local.json_path_parser import JsonPathItem, JsonPathItemType
def set_value_by_key(payload: dict, keys: List[JsonPathItem], value: Any) -> None:
"""
Set value in payload by key.
Args:
payload: arbitrary json-like object
keys:
list of json path items, e.g.:
[
JsonPathItem(item_type=<JsonPathItemType.KEY: 'key'>, value='a'),
JsonPathItem(item_type=<JsonPathItemType.INDEX: 'index'>, value=0),
JsonPathItem(item_type=<JsonPathItemType.INDEX: 'index'>, value=1),
JsonPathItem(item_type=<JsonPathItemType.KEY: 'key'>, value='b')
]
The original keys could look like this:
- "name"
- "address.city"
- "location[].name"
- "location[0].name"
value: value to set
"""
Setter.set(payload, keys.copy(), value, None, None)
class Setter:
TYPE: Any
SETTERS: Dict[JsonPathItemType, Type["Setter"]] = {}
@classmethod
def add_setter(cls, item_type: JsonPathItemType, setter: Type["Setter"]) -> None:
cls.SETTERS[item_type] = setter
@classmethod
def set(
cls,
data: Any,
k_list: List[JsonPathItem],
value: Dict[str, Any],
prev_data: Any,
prev_key: Optional[JsonPathItem],
) -> None:
if not k_list:
return
current_key = k_list.pop(0)
cls.SETTERS[current_key.item_type]._set(
data,
current_key,
k_list,
value,
prev_data,
prev_key,
)
@classmethod
def _set(
cls,
data: Any,
current_key: JsonPathItem,
k_list: List[JsonPathItem],
value: Dict[str, Any],
prev_data: Any,
prev_key: Optional[JsonPathItem],
) -> None:
if isinstance(data, cls.TYPE):
cls._set_compatible_types(
data=data, current_key=current_key, k_list=k_list, value=value
)
else:
cls._set_incompatible_types(
current_key=current_key,
k_list=k_list,
value=value,
prev_data=prev_data,
prev_key=prev_key,
)
@classmethod
def _set_compatible_types(
cls,
data: Any,
current_key: JsonPathItem,
k_list: List[JsonPathItem],
value: Dict[str, Any],
) -> None:
raise NotImplementedError()
@classmethod
def _set_incompatible_types(
cls,
current_key: JsonPathItem,
k_list: List[JsonPathItem],
value: Dict[str, Any],
prev_data: Any,
prev_key: Optional[JsonPathItem],
) -> None:
raise NotImplementedError()
class KeySetter(Setter):
TYPE = Dict
@classmethod
def _set_compatible_types(
cls,
data: Any,
current_key: JsonPathItem,
k_list: List[JsonPathItem],
value: Dict[str, Any],
) -> None:
if current_key.key not in data:
data[current_key.key] = {}
if len(k_list) == 0:
if isinstance(data[current_key.key], dict):
data[current_key.key].update(value)
else:
data[current_key.key] = value
else:
cls.set(data[current_key.key], k_list.copy(), value, data, current_key)
@classmethod
def _set_incompatible_types(
cls,
current_key: JsonPathItem,
k_list: List[JsonPathItem],
value: Dict[str, Any],
prev_data: Any,
prev_key: Optional[JsonPathItem],
) -> None:
assert prev_key is not None
if len(k_list) == 0:
if prev_key.item_type == JsonPathItemType.KEY:
prev_data[prev_key.key] = {current_key.key: value}
else: # if prev key was WILDCARD, we need to pass INDEX instead with an index set
prev_data[prev_key.index] = {current_key.key: value}
else:
if prev_key.item_type == JsonPathItemType.KEY:
prev_data[prev_key.key] = {current_key.key: {}}
cls.set(
prev_data[prev_key.key][current_key.key],
k_list.copy(),
value,
prev_data[prev_key.key],
current_key,
)
else:
prev_data[prev_key.index] = {current_key.key: {}}
cls.set(
prev_data[prev_key.index][current_key.key],
k_list.copy(),
value,
prev_data[prev_key.index],
current_key,
)
class _ListSetter(Setter):
TYPE = List
@classmethod
def _set_incompatible_types(
cls,
current_key: JsonPathItem,
k_list: List[JsonPathItem],
value: Dict[str, Any],
prev_data: Any,
prev_key: Optional[JsonPathItem],
) -> None:
assert prev_key is not None
if prev_key.item_type == JsonPathItemType.KEY:
prev_data[prev_key.key] = []
return
else:
prev_data[prev_key.index] = []
return
@classmethod
def _set_compatible_types(
cls,
data: Any,
current_key: JsonPathItem,
k_list: List[JsonPathItem],
value: Dict[str, Any],
) -> None:
raise NotImplementedError()
class IndexSetter(_ListSetter):
@classmethod
def _set_compatible_types(
cls,
data: Any,
current_key: JsonPathItem,
k_list: List[JsonPathItem],
value: Dict[str, Any],
) -> None:
assert current_key.index is not None
if current_key.index < len(data):
if len(k_list) == 0:
if isinstance(data[current_key.index], dict):
data[current_key.index].update(value)
else:
data[current_key.index] = value
return
cls.set(data[current_key.index], k_list.copy(), value, data, current_key)
class WildcardIndexSetter(_ListSetter):
@classmethod
def _set_compatible_types(
cls,
data: Any,
current_key: JsonPathItem,
k_list: List[JsonPathItem],
value: Dict[str, Any],
) -> None:
if len(k_list) == 0:
for i, item in enumerate(data):
if isinstance(item, dict):
data[i].update(value)
else:
data[i] = value
else:
for i, item in enumerate(data):
cls.set(
item,
k_list.copy(),
value,
data,
JsonPathItem(item_type=JsonPathItemType.INDEX, index=i),
)
Setter.add_setter(JsonPathItemType.KEY, KeySetter)
Setter.add_setter(JsonPathItemType.INDEX, IndexSetter)
Setter.add_setter(JsonPathItemType.WILDCARD_INDEX, WildcardIndexSetter)