platypush/platypush/schemas/system/_connection/_base.py
Fabio Manganiello ebe79ac29a
Refactored system schema dataclasses.
- `percent_field` should be declared on `platypush.schemas.dataclasses`
  level, since it's not specific to the `system` plugin.
- Added a common `SystemBaseSchema` that takes care of calling
  `_asdict()` if the object is passed as a `psutil` object instead of a
  dict.
2023-04-22 17:19:23 +02:00

39 lines
1.1 KiB
Python

from marshmallow import pre_load
from .._base import SystemBaseSchema
class ConnectionBaseSchema(SystemBaseSchema):
"""
Base schema for connections.
"""
@pre_load
def pre_load(self, data, **_) -> dict:
data = super().pre_load(data)
addr_mapping = {
'laddr': ('local_address', 'local_port'),
'raddr': ('remote_address', 'remote_port'),
}
# Parse laddr/raddr attributes
for ext_attr, (addr_attr, port_attr) in addr_mapping.items():
value = data.pop(ext_attr, None)
if not value:
data[addr_attr] = data[port_attr] = None
elif isinstance(value, tuple):
data[addr_attr], data[port_attr] = value
elif isinstance(value, str):
data[addr_attr] = value
data[port_attr] = None
# Handle enum values
for attr in ['type', 'family']:
value = data.pop(attr, None)
if value is not None:
data[attr] = value.name
if data.get('status') == 'NONE':
data['status'] = None
return data