fix(base_quantity): Removed string return from serialization of pint quantity.

src/infrasys/base_quantity.py

@@ -30,26 +30,24 @@ def __init_subclass__(cls, **kwargs):
     # Required for pydantic validation
     @classmethod
     def __get_pydantic_core_schema__(
-        cls, _: Any, handler: GetCoreSchemaHandler
+        cls, _: Type["BaseQuantity"], handler: GetCoreSchemaHandler
     ) -> core_schema.CoreSchema:
         return core_schema.with_info_after_validator_function(
             cls._validate,
-            core_schema.union_schema(
-                [
-                    handler(pint.Quantity),
-                    core_schema.float_schema(),
-                ]
-            ),
+            core_schema.any_schema(),
             field_name=handler.field_name,
             serialization=core_schema.plain_serializer_function_ser_schema(
-                cls._serialize, info_arg=True, return_schema=core_schema.str_schema()
+                cls._serialize,
+                info_arg=True,
+                return_schema=core_schema.any_schema(),
             ),
         )
 
     # Required for pydantic validation
     @classmethod
     def _validate(cls, field_value: Any, _: core_schema.ValidationInfo) -> "BaseQuantity":
-        if isinstance(field_value, BaseQuantity):
+        # Type check is more robubst to check that is not an instance of a bare "BaseQuantity"
+        if type(field_value) is cls:
             if cls.__base_unit__:
                 assert field_value.check(
                     cls.__base_unit__
@@ -61,22 +59,19 @@ def _validate(cls, field_value: Any, _: core_schema.ValidationInfo) -> "BaseQuan
                     cls.__base_unit__
                 ), f"Unit must be compatible with {cls.__base_unit__}"
                 return cls(field_value.magnitude, field_value.units)
-            else:
-                raise ValueError(f"Invalid type for BaseQuantity: {type(field_value)}")
-        if isinstance(field_value, cls):
-            return field_value
-        if isinstance(field_value, float) or isinstance(field_value, int):
-            return cls(field_value, cls.__base_unit__)
-        raise TypeError("Type not supported")
+        return cls(field_value, cls.__base_unit__)
 
     @classmethod
-    def _serialize(cls, input_value, info: SerializationInfo):
-        if context := info.context:
+    def _serialize(cls, input_value, info: SerializationInfo) -> float | str:
+        return_value = input_value
+        if context := info.context:  # type: ignore
             # We can add more logic that will change the serialization here.
             magnitude_only = context.get("magnitude_only")
             if magnitude_only:
-                return str(input_value.magnitude)
-        return str(input_value)
+                return return_value.magnitude
+        if info.mode == "json":
+            return_value = str(return_value)
+        return return_value
 
     def to_dict(self) -> dict[str, Any]:
         """Convert a quantity to a dictionary for serialization."""

tests/test_base_quantity.py

@@ -1,5 +1,4 @@
 from pydantic import ValidationError
-import pydantic_core
 from infrasys.base_quantity import ureg, BaseQuantity
 from infrasys.component import Component
 from infrasys.quantities import ActivePower, Time, Voltage
@@ -76,8 +75,10 @@ class _(BaseQuantity):
     with pytest.raises(ValidationError):
         BaseQuantityComponent(name="test", voltage=Voltage(test_magnitude, "meter"))
 
-    with pytest.raises(pydantic_core.ValidationError):
-        BaseQuantityComponent(name="test", voltage=[0, 1])
+    test_component = BaseQuantityComponent(name="test", voltage=[0, 1])
+    assert type(test_component.voltage) is Voltage
+    assert test_component.voltage.magnitude.tolist() == [0, 1]
+    assert test_component.voltage.units == test_unit
 
 
 @pytest.mark.parametrize("input_unit", [Voltage(10, "kV"), 10 * ureg.volt, 10, 10.0])
@@ -95,5 +96,8 @@ def test_custom_serialization():
 
     assert model_dump["voltage"] == str(Voltage(10.0, "volt"))
 
+    model_dump = component.model_dump(context={"magnitude_only": True})
+    assert model_dump["voltage"] == 10.0
+
     model_dump = component.model_dump(mode="json", context={"magnitude_only": True})
-    assert model_dump["voltage"] == str(10.0)
+    assert model_dump["voltage"] == 10.0