Dataset.__setitem__ now converts Models as input using to_data() + from_data()

Author: sveinugu

src/omnipy/data/dataset.py

@@ -227,19 +227,24 @@ def __init__(  # noqa: C901
         self._init(super_kwargs, **kwargs)
 
         try:
-            super().__init__(**super_kwargs)
+            self._primary_validation(super_kwargs)
         except ValidationError:
             if model_or_dataset_as_input:
-                super().__init__()
-                self.from_data(super_kwargs[DATA_KEY])
+                self._secondary_validation_from_data(super_kwargs)
             else:
                 raise
 
-        UserDict.__init__(self, self.data)  # noqa
-
         if not self.__doc__:
             self._set_standard_field_description()
 
+    def _primary_validation(self, super_kwargs):
+        # Pydantic validation of super_kwargs
+        super().__init__(**super_kwargs)
+
+    def _secondary_validation_from_data(self, super_kwargs):
+        super().__init__()
+        self.from_data(super_kwargs[DATA_KEY])
+
     def _init(self, super_kwargs: dict[str, Any], **kwargs: Any) -> None:
         ...
 
@@ -390,7 +395,7 @@ def _set_data_file_and_validate(self, key: str, val: ModelT) -> None:
 
         try:
             self.data[key] = val
-            self._validate(key)
+            self._validate_data_file(key)
         except Exception:
             if has_prev_value:
                 self.data[key] = prev_value
@@ -408,7 +413,14 @@ def update_forward_refs(cls, **localns: Any) -> None:
         cls.__name__ = remove_forward_ref_notation(cls.__name__)
         cls.__qualname__ = remove_forward_ref_notation(cls.__qualname__)
 
-    def _validate(self, _data_file: str) -> None:
+    def _validate_data_file(self, _data_file: str) -> None:
+        val = self.data[_data_file]
+        if is_model_instance(val):
+            self.data[_data_file] = self.get_model_class()(val)
+        else:
+            self._force_full_validation()
+
+    def _force_full_validation(self):
         self.data = self.data  # Triggers pydantic validation, as validate_assignment=True
 
     def __iter__(self) -> Iterator:
@@ -461,7 +473,7 @@ def _from_dict_with_callback(self,
         for data_file, contents in data.items():
             new_model = model_cls()
             callback_func(new_model, contents)
-            self[data_file] = new_model
+            self.data[data_file] = new_model
 
     def absorb(self, other: 'Dataset'):
         self.from_data(other.to_data(), update=True)
@@ -706,7 +718,7 @@ def set_model(self, data_file: str, model: GeneralModelT) -> None:
         try:
             self._custom_field_models[data_file] = model
             if data_file in self.data:
-                self._validate(data_file)
+                self._validate_data_file(data_file)
             else:
                 self.data[data_file] = model()
         except ValidationError:
@@ -719,15 +731,26 @@ def get_model(self, data_file: str) -> GeneralModelT:
         else:
             return self.get_model_class()
 
-    def _validate(self, data_file: str) -> None:
+    def from_data(self,
+                  data: dict[str, Any] | Iterator[tuple[str, Any]],
+                  update: bool = True) -> None:
+        super().from_data(data, update)
+        for data_file in self:
+            self._validate_data_file_according_to_custom_field_model(data_file)
+        self._force_full_validation()
+
+    def _validate_data_file(self, data_file: str) -> None:
+        self._validate_data_file_according_to_custom_field_model(data_file)
+        self._force_full_validation()
+
+    def _validate_data_file_according_to_custom_field_model(self, data_file: str):
         if data_file in self._custom_field_models:
             model = self._custom_field_models[data_file]
             if not is_model_instance(model):
                 model = Model[model]
             data_obj = self._to_data_if_model(self.data[data_file])
             parsed_data = self._to_data_if_model(model(data_obj))
             self.data[data_file] = parsed_data
-        super()._validate(data_file)  # validates all data according to ModelNewT
 
     @staticmethod
     def _to_data_if_model(data_obj: Any):

tests/data/test_dataset.py

@@ -20,7 +20,7 @@
                                MyFwdRefDataset,
                                MyNestedFwdRefDataset,
                                ParamUpperStrDataset)
-from .helpers.models import NumberModel, StringToLength
+from .helpers.models import MyFloatObjModel, NumberModel, StringToLength
 
 
 def test_no_model():
@@ -435,6 +435,19 @@ def test_set_items_with_tuple_or_list() -> None:
     assert dataset['_untitled_4'] == dataset[-1] == Model[int](999)
 
 
+def test_set_item_converting_models_as_input():
+    float_model = Model[float](4.5)
+    my_float_model = MyFloatObjModel(MyFloatObject(int_part=4, float_part=0.5))
+
+    float_dataset = Dataset[Model[float]]()
+    float_dataset['x'] = my_float_model
+    assert float_dataset['x'].contents == 4.5
+
+    my_float_dataset = MyFloatObjDataset()
+    my_float_dataset['x'] = float_model
+    assert my_float_dataset['x'].contents == MyFloatObject(int_part=4, float_part=0.5)
+
+
 def test_del_item_with_str() -> None:
     dataset = Dataset[Model[int]](data_file_1=123, data_file_2=456, data_file_3=789)
 

tests/modules/remote/test_tasks.py

@@ -152,16 +152,10 @@ def _assert_json_query_results(assert_model_if_dyn_conv_else_val, data: JsonData
     assert_model_if_dyn_conv_else_val(data['delillos_2']['lyrics'][0], str, 'Joda sier Arne')
 
 
-# TODO: Models as input to Dataset.__setitem__ is not converted to the model of the dataset in the
-#       same way as in the __init__method, causing this test to fail if
-#       `dynamically_convert_elements_to_models == True`
-
-
 @pc.parametrize_with_cases('case', cases='.cases.request_types')
 async def test_get_from_api_endpoint_without_session(
     query_urls: Annotated[HttpUrlDataset, pytest.fixture],
     assert_model_if_dyn_conv_else_val: Annotated[AssertModelOrValFunc, pytest.fixture],
-    skip_test_if_dynamically_convert_elements_to_models: Annotated[None, pytest.fixture],
     case: RequestTypeCase,
 ) -> None:
     data = await case.job.run(query_urls, **case.kwargs)