Implemented flexible __setitem__ of Dataset

Author: sveinugu

src/omnipy/data/dataset.py

@@ -4,7 +4,7 @@
 import os
 import tarfile
 from tempfile import TemporaryDirectory
-from typing import Any, Callable, cast, Generic, Iterator
+from typing import Any, Callable, cast, Generic, Iterator, MutableMapping
 from urllib.parse import ParseResult, urlparse
 
 import humanize
@@ -23,6 +23,12 @@
                                  is_model_instance,
                                  waiting_for_terminal_repr)
 from omnipy.data.model import Model
+from omnipy.data.selector import (create_updated_mapping,
+                                  Index2DataItemsType,
+                                  Key2DataItemType,
+                                  prepare_selected_items_with_iterable_data,
+                                  prepare_selected_items_with_mapping_data,
+                                  select_keys)
 from omnipy.util.helpers import (get_calling_module_name,
                                  get_default_if_typevar,
                                  is_iterable,
@@ -272,40 +278,73 @@ def _get_standard_field_description(cls) -> str:
                 'particular specialization of the Model class. Both main classes are wrapping '
                 'the excellent Python package named `pydantic`.')
 
-    def __setitem__(self, data_file: str, data_obj: Any) -> None:
-        has_prev_value = data_file in self.data
-        prev_value = self.data.get(data_file)
+    def __getitem__(self, selector: str | int | slice | Iterable[str | int]) -> Any:
+        selected_keys = select_keys(selector, self.data)
+
+        if selected_keys.singular:
+            return self.data[selected_keys.keys[0]]
+        else:
+            return self.__class__({key: self.data[key] for key in selected_keys.keys})
+
+    def __setitem__(
+        self,
+        selector: str | int | slice | Iterable[str | int],
+        data_obj: dict[str, ModelT] | Iterable[ModelT] | ModelT,
+    ) -> None:
+        selected_keys = select_keys(selector, self.data)
+
+        if selected_keys.singular:
+            self._set_data_file_and_validate(selected_keys.keys[0], cast(ModelT, data_obj))
+        else:
+            key_2_data_item: Key2DataItemType[ModelT]
+            index_2_data_items: Index2DataItemsType[ModelT]
+
+            if isinstance(data_obj, MutableMapping):
+                key_2_data_item, index_2_data_items = \
+                    prepare_selected_items_with_mapping_data(
+                        selected_keys.keys, selected_keys.last_index, data_obj,)
+
+            elif is_iterable(data_obj):
+                key_2_data_item, index_2_data_items = \
+                    prepare_selected_items_with_iterable_data(
+                        selected_keys.keys, selected_keys.last_index, tuple(data_obj), self.data)
 
-        try:
-            self.data[data_file] = data_obj
-            self._validate(data_file)
-        except:  # noqa
-            if has_prev_value:
-                self.data[data_file] = prev_value
             else:
-                del self.data[data_file]
+                raise TypeError('Data object must be a mapping or an iterable')
+
+            self._update_selected_items_with_data_items(key_2_data_item, index_2_data_items)
+
+    def _update_selected_items_with_data_items(
+        self,
+        key_2_data_item: Key2DataItemType[ModelT],
+        index_2_data_item: Index2DataItemsType[ModelT],
+    ) -> None:
+
+        updated_mapping = create_updated_mapping(self.data, key_2_data_item, index_2_data_item)
+        self._replace_data_with_mapping(updated_mapping)
+
+    def _replace_data_with_mapping(self, updated_mapping):
+        prev_data = self.data
+        try:
+            self.absorb_and_replace(self.__class__(updated_mapping))
+        except Exception:
+            self.data = prev_data
             raise
 
-    def __getitem__(self, selector: str | int | slice | Iterable[str | int]) -> Any:
-        if isinstance(selector, str):
-            if selector in self.data:
-                return self.data[selector]
-            else:
-                return self.data[selector]
-        else:
-            data_keys = tuple(self.data.keys())
-
-            if isinstance(selector, int):
-                return self.data[data_keys[selector]]
-            elif isinstance(selector, slice):
-                return self.__class__({key: self.data[key] for key in data_keys[selector]})
-            elif is_iterable(selector):
-                selected_keys = [data_keys[_] if isinstance(_, int) else _ for _ in selector]
-                return self.__class__({key: self.data[key] for key in selected_keys})
+    def _set_data_file_and_validate(self, key: str, val: ModelT) -> None:
+        has_prev_value = key in self.data
+        if has_prev_value:
+            prev_value = self.data[key]
+
+        try:
+            self.data[key] = val
+            self._validate(key)
+        except Exception:
+            if has_prev_value:
+                self.data[key] = prev_value
             else:
-                raise KeyError(
-                    'Selector is of incorrect type. Must be a string, a positive integer,'
-                    'or a slice (e.g. `dataset[2:5]`).')
+                del self.data[key]
+            raise
 
     @classmethod
     def update_forward_refs(cls, **localns: Any) -> None:

src/omnipy/data/selector.py

@@ -0,0 +1,167 @@
+from collections import defaultdict
+from dataclasses import dataclass
+from typing import Iterable, Mapping, MutableMapping, TypeAlias
+
+from typing_extensions import TypeVar
+
+from omnipy.util.helpers import is_iterable
+
+T = TypeVar('T')
+
+UNTITLED_KEY = '_untitled'
+
+
+@dataclass
+class SelectedKeys:
+    singular: bool
+    keys: tuple[str, ...]
+    last_index: int = -1
+
+
+Key2DataItemType: TypeAlias = dict[str, tuple[str, T] | None]
+Index2DataItemsType: TypeAlias = defaultdict[int, list[tuple[str, T]]]
+MappingType: TypeAlias = MutableMapping[str, T]
+
+
+def select_keys(selector: str | int | slice | Iterable[str | int],
+                mapping: MappingType[T]) -> SelectedKeys:
+    if isinstance(selector, str):
+        return SelectedKeys(singular=True, keys=(selector,))
+    else:
+        data_keys = tuple(mapping.keys())
+
+        if isinstance(selector, int):
+            return SelectedKeys(singular=True, keys=(data_keys[selector],))
+
+        if isinstance(selector, slice):
+            last_index = selector.indices(len(data_keys))[1] - 1
+            return SelectedKeys(singular=False, keys=data_keys[selector], last_index=last_index)
+
+        elif is_iterable(selector):
+            keys = tuple(data_keys[_] if isinstance(_, int) else _ for _ in selector)
+            if keys and keys[-1] in data_keys:
+                last_index = data_keys.index(keys[-1])
+            else:
+                last_index = len(data_keys) - 1
+            return SelectedKeys(singular=False, keys=keys, last_index=last_index)
+
+        else:
+            raise KeyError('Selector is of incorrect type. Must be a string, a positive integer,'
+                           'or a slice (e.g. `dataset[2:5]`).')
+
+
+def prepare_selected_items_with_mapping_data(
+    keys: tuple[str, ...],
+    last_index: int,
+    data_obj: MappingType[T],
+) -> tuple[Key2DataItemType[T], Index2DataItemsType[T]]:
+
+    data_obj_keys = tuple(data_obj.keys())
+    key_2_data_item: Key2DataItemType[T] = {}
+    index_2_data_items: Index2DataItemsType[T] = defaultdict(list)
+
+    for i, data_key in enumerate(data_obj.keys()):
+        if i < len(keys):
+            key_2_data_item[keys[i]] = (data_key, data_obj[data_key])
+        else:
+            index_2_data_items[last_index].extend((key, data_obj[key]) for key in data_obj_keys[i:])
+            break
+
+    if len(keys) > len(data_obj_keys):
+        for key in keys[len(data_obj_keys):]:
+            key_2_data_item[key] = None
+
+    return key_2_data_item, index_2_data_items
+
+
+def prepare_selected_items_with_iterable_data(
+    keys: tuple[str, ...],
+    last_index: int,
+    data_obj: tuple[T, ...],
+    mapping: MappingType[T],
+) -> tuple[Key2DataItemType[T], Index2DataItemsType[T]]:
+
+    key_2_data_item: Key2DataItemType[T] = {}
+    index_2_data_items: Index2DataItemsType[T] = defaultdict(list)
+
+    for i, data_val in enumerate(data_obj):
+        if i < len(keys):
+            if keys[i] in mapping:
+                key_2_data_item[keys[i]] = (keys[i], data_val)
+            else:
+                index_2_data_items[last_index].append((keys[i], data_val))
+        else:
+            index_2_data_items[last_index].extend(
+                (UNTITLED_KEY, val) for j, val in enumerate(data_obj[i:]))
+            break
+
+    if len(keys) > len(data_obj):
+        for key in keys[len(data_obj):]:
+            key_2_data_item[key] = None
+
+    return key_2_data_item, index_2_data_items
+
+
+def create_updated_mapping(
+    mapping: MappingType[T],
+    key_2_data_item: Key2DataItemType[T],
+    index_2_data_item: Index2DataItemsType[T],
+) -> MappingType[T]:
+
+    updated_mapping: dict[str, T] = {}
+
+    uniquely_add_extra_items_by_index_to_mapping(-1, index_2_data_item, updated_mapping)
+
+    for i, (key, val) in enumerate(mapping.items()):
+        if key in key_2_data_item:
+            uniquely_add_item_by_key_to_mapping_if_val(key, key_2_data_item, updated_mapping)
+        else:
+            uniquely_add_item_to_mapping(key, val, updated_mapping)
+
+        uniquely_add_extra_items_by_index_to_mapping(i, index_2_data_item, updated_mapping)
+
+    return updated_mapping
+
+
+def uniquely_add_extra_items_by_index_to_mapping(
+    index: int,
+    index_2_data_item: Index2DataItemsType[T],
+    mapping: MappingType[T],
+) -> None:
+    if index in index_2_data_item:
+        for key, val in index_2_data_item[index]:
+            uniquely_add_item_to_mapping(key, val, mapping)
+
+
+def uniquely_add_item_by_key_to_mapping_if_val(
+    key: str,
+    key_2_data_item: Key2DataItemType[T],
+    mapping: MappingType[T],
+):
+    data_item = key_2_data_item[key]
+    if data_item is not None:
+        key, val = data_item
+        uniquely_add_item_to_mapping(key, val, mapping)
+
+
+def uniquely_add_item_to_mapping(key: str, val: T, mapping: MappingType[T]) -> None:
+    mapping[make_unique_key(key, mapping)] = val
+
+
+def make_unique_key(key: str, mapping: MappingType[T]) -> str:
+    while key in mapping:
+        if is_duplicate_name(key):
+            key = increase_duplicate_count(key)
+        else:
+            key = f'{key}_2'
+    return key
+
+
+def is_duplicate_name(key: str) -> bool:
+    splitted_key = key.rsplit('_', 1)
+    return len(splitted_key) == 2 and splitted_key[1].isdigit()
+
+
+def increase_duplicate_count(key: str) -> str:
+    key, count = key.rsplit('_', 1)
+    return f'{key}_{int(count) + 1}'