1414 even in minor or patch releases.
1515"""
1616
17+ from contextvars import ContextVar
1718from typing import Any , Type
1819
19- from marshmallow import Schema , ValidationError , fields
20+ from marshmallow import Schema , ValidationError
21+ from marshmallow .fields import Field
2022
2123from .._apparent_power import ApparentPower
2224from .._current import Current
2931from .._temperature import Temperature
3032from .._voltage import Voltage
3133
34+ serialize_as_string_default : ContextVar [bool ] = ContextVar (
35+ "serialize_as_string_default" , default = False
36+ )
37+ """Context variable to control the default serialization format for quantities.
3238
33- class _QuantityField (fields .Field ):
39+ If True, quantities are serialized as strings with units.
40+ If False, quantities are serialized as floats.
41+
42+ This can be overridden on a per-field basis using the `serialize_as_string`
43+ metadata attribute.
44+ """
45+
46+
47+ class _QuantityField (Field [Any ]):
3448 """Custom field for Quantity objects supporting per-field serialization configuration.
3549
3650 This class handles serialization and deserialization of ALL Quantity
@@ -57,24 +71,34 @@ class _QuantityField(fields.Field):
5771 field_type : Type [Quantity ] | None = None
5872 """The specific Quantity subclass."""
5973
74+ def __init__ (self , * args : Any , ** kwargs : Any ) -> None :
75+ """Initialize the field."""
76+ self .serialize_as_string_override = kwargs .pop ("serialize_as_string" , None )
77+ super ().__init__ (* args , ** kwargs )
78+
6079 def _serialize (
61- self , value : Quantity , attr : str | None , obj : Any , ** kwargs : Any
80+ self , value : Any | None , attr : str | None , obj : Any , ** kwargs : Any
6281 ) -> Any :
6382 """Serialize the Quantity object based on per-field configuration."""
6483 if self .field_type is None or not issubclass (self .field_type , Quantity ):
6584 raise TypeError (
6685 "field_type must be set to a Quantity subclass in the subclass."
6786 )
87+ if value is None :
88+ return None
6889
69- assert self .parent is not None
90+ if not isinstance (value , Quantity ):
91+ raise TypeError (
92+ f"Expected a Quantity object, but got { type (value ).__name__ } ."
93+ )
7094
7195 # Determine the serialization format
72- default = (
73- False
74- if self .parent .context is None
75- else self .parent .context .get ("serialize_as_string_default" , False )
96+ default = serialize_as_string_default .get ()
97+ serialize_as_string = (
98+ self .serialize_as_string_override
99+ if self .serialize_as_string_override is not None
100+ else default
76101 )
77- serialize_as_string = self .metadata .get ("serialize_as_string" , default )
78102
79103 if serialize_as_string :
80104 # Use the Quantity's native string representation (includes unit)
@@ -177,7 +201,7 @@ class VoltageField(_QuantityField):
177201 field_type = Voltage
178202
179203
180- QUANTITY_FIELD_CLASSES : dict [type [Quantity ], type [fields . Field ]] = {
204+ QUANTITY_FIELD_CLASSES : dict [type [Quantity ], type [Field [ Any ] ]] = {
181205 ApparentPower : ApparentPowerField ,
182206 Current : CurrentField ,
183207 Energy : EnergyField ,
@@ -208,8 +232,10 @@ class QuantitySchema(Schema):
208232 from marshmallow_dataclass import class_schema
209233 from marshmallow.validate import Range
210234 from frequenz.quantities import Percentage
211- from frequenz.quantities.experimental.marshmallow import QuantitySchema
212- from typing import cast
235+ from frequenz.quantities.experimental.marshmallow import (
236+ QuantitySchema,
237+ serialize_as_string_default,
238+ )
213239
214240 @dataclass
215241 class Config:
@@ -245,29 +271,24 @@ class Config:
245271 },
246272 )
247273
248- @classmethod
249- def load(cls, config: dict[str, Any]) -> "Config":
250- schema = class_schema(cls, base_schema=QuantitySchema)(
251- serialize_as_string_default=True # type: ignore[call-arg]
252- )
253- return cast(Config, schema.load(config))
274+ config_obj = Config()
275+ Schema = class_schema(Config, base_schema=QuantitySchema)
276+ schema = Schema()
277+
278+ # Default serialization (as float)
279+ result = schema.dump(config_obj)
280+ assert result["percentage_serialized_as_schema_default"] == 25.0
281+
282+ # Override default serialization to string
283+ serialize_as_string_default.set(True)
284+ result = schema.dump(config_obj)
285+ assert result["percentage_serialized_as_schema_default"] == "25.0 %"
286+ serialize_as_string_default.set(False) # Reset context
287+
288+ # Per-field configuration always takes precedence
289+ assert result["percentage_always_as_string"] == "25.0 %"
290+ assert result["percentage_always_as_float"] == 25.0
254291 ```
255292 """
256293
257- TYPE_MAPPING : dict [type [Quantity ], type [fields .Field ]] = QUANTITY_FIELD_CLASSES
258-
259- def __init__ (
260- self , * args : Any , serialize_as_string_default : bool = False , ** kwargs : Any
261- ) -> None :
262- """
263- Initialize the schema with a default serialization format.
264-
265- Args:
266- *args: Additional positional arguments.
267- serialize_as_string_default: Default serialization format for quantities.
268- If True, quantities are serialized as strings with units.
269- If False, quantities are serialized as floats.
270- **kwargs: Additional keyword arguments.
271- """
272- super ().__init__ (* args , ** kwargs )
273- self .context ["serialize_as_string_default" ] = serialize_as_string_default
294+ TYPE_MAPPING : dict [type , type [Field [Any ]]] = QUANTITY_FIELD_CLASSES
0 commit comments