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.
This commit is contained in:
Fabio Manganiello 2023-04-22 13:11:48 +02:00
parent 2d618188c8
commit ebe79ac29a
Signed by: blacklight
GPG key ID: D90FBA7F76362774
12 changed files with 61 additions and 41 deletions

View file

@ -151,7 +151,7 @@ class SystemPlugin(SensorPlugin, EntityManager):
def _cpu_stats(self) -> CpuStats:
stats = psutil.cpu_stats()
return CpuStatsSchema().load(stats._asdict()) # type: ignore
return CpuStatsSchema().load(stats) # type: ignore
@action
def cpu_stats(self) -> CpuStats:
@ -164,11 +164,11 @@ class SystemPlugin(SensorPlugin, EntityManager):
def _cpu_frequency_avg(self) -> CpuFrequency:
freq = psutil.cpu_freq(percpu=False)
return CpuFrequencySchema().load(freq._asdict()) # type: ignore
return CpuFrequencySchema().load(freq) # type: ignore
def _cpu_frequency_per_cpu(self) -> List[CpuFrequency]:
freq = psutil.cpu_freq(percpu=True)
return CpuFrequencySchema().load(freq._asdict(), many=True) # type: ignore
return CpuFrequencySchema().load(freq, many=True) # type: ignore
@action
def cpu_frequency(
@ -198,9 +198,7 @@ class SystemPlugin(SensorPlugin, EntityManager):
return psutil.getloadavg()
def _mem_virtual(self) -> MemoryStats:
return MemoryStatsSchema().load(
psutil.virtual_memory()._asdict()
) # type: ignore
return MemoryStatsSchema().load(psutil.virtual_memory()) # type: ignore
@action
def mem_virtual(self) -> dict:
@ -212,7 +210,7 @@ class SystemPlugin(SensorPlugin, EntityManager):
return MemoryStatsSchema().dump(self._mem_virtual()) # type: ignore
def _mem_swap(self) -> SwapStats:
return SwapStatsSchema().load(psutil.swap_memory()._asdict()) # type: ignore
return SwapStatsSchema().load(psutil.swap_memory()) # type: ignore
@action
def mem_swap(self) -> dict:
@ -226,7 +224,6 @@ class SystemPlugin(SensorPlugin, EntityManager):
def _disk_info(self) -> List[Disk]:
parts = {part.device: part._asdict() for part in psutil.disk_partitions()}
basename_parts = {os.path.basename(part): part for part in parts}
io_stats = {
basename_parts[disk]: stats._asdict()
for disk, stats in psutil.disk_io_counters(perdisk=True).items()
@ -234,7 +231,7 @@ class SystemPlugin(SensorPlugin, EntityManager):
}
usage = {
disk: psutil.disk_usage(info['mountpoint'])._asdict()
disk: psutil.disk_usage(info['mountpoint'])._asdict() # type: ignore
for disk, info in parts.items()
}

View file

@ -1,3 +1,4 @@
from dataclasses import field
from datetime import date, datetime
from uuid import UUID
@ -6,11 +7,26 @@ from marshmallow import (
Schema,
fields,
post_dump,
pre_load,
)
from marshmallow.validate import Range
from .. import Date, DateTime
def percent_field(**kwargs):
"""
Field used to model percentage float fields between 0 and 1.
"""
return field(
default_factory=float,
metadata={
'validate': Range(min=0, max=1),
**kwargs,
},
)
class DataClassSchema(Schema):
"""
Base schema class for data classes that support Marshmallow schemas.
@ -45,6 +61,10 @@ class DataClassSchema(Schema):
return matching_fields[0]
@pre_load
def pre_load(self, data, **__) -> dict:
return data
@post_dump
def post_dump(self, data: dict, **__) -> dict:
# Use data_key parameters only for load

View file

@ -1,16 +1,15 @@
from dataclasses import field
from marshmallow import pre_load
from marshmallow.validate import Range
from platypush.schemas.dataclasses import DataClassSchema
def percent_field(**kwargs):
class SystemBaseSchema(DataClassSchema):
"""
Field used to model percentage float fields between 0 and 1.
Base schema for system info.
"""
return field(
default_factory=float,
metadata={
'validate': Range(min=0, max=1),
**kwargs,
},
)
@pre_load
def pre_load(self, data, **_) -> dict:
if hasattr(data, '_asdict'):
data = data._asdict()
return data

View file

@ -1,18 +1,16 @@
from marshmallow import pre_load
from platypush.schemas.dataclasses import DataClassSchema
from .._base import SystemBaseSchema
class ConnectionBaseSchema(DataClassSchema):
class ConnectionBaseSchema(SystemBaseSchema):
"""
Base schema for connections.
"""
@pre_load
def pre_load(self, data, **_) -> dict:
if hasattr(data, '_asdict'):
data = data._asdict()
data = super().pre_load(data)
addr_mapping = {
'laddr': ('local_address', 'local_port'),
'raddr': ('remote_address', 'remote_port'),

View file

@ -1,15 +1,16 @@
from marshmallow import pre_load
from platypush.schemas.dataclasses import DataClassSchema
from .._base import SystemBaseSchema
class CpuInfoBaseSchema(DataClassSchema):
class CpuInfoBaseSchema(SystemBaseSchema):
"""
Base schema for CPU info.
"""
@pre_load
def pre_load(self, data: dict, **_) -> dict:
data = super().pre_load(data)
if data.get('hz_advertised'):
data['frequency_advertised'] = data.pop('hz_advertised')[0]
if data.get('hz_actual'):
@ -18,7 +19,7 @@ class CpuInfoBaseSchema(DataClassSchema):
return data
class CpuTimesBaseSchema(DataClassSchema):
class CpuTimesBaseSchema(SystemBaseSchema):
"""
Base schema for CPU times.
"""
@ -29,6 +30,7 @@ class CpuTimesBaseSchema(DataClassSchema):
Convert the underlying object to dict and normalize all the percentage
values from [0, 100] to [0, 1].
"""
data = super().pre_load(data)
return {
key: value / 100.0
for key, value in (

View file

@ -1,7 +1,7 @@
from dataclasses import dataclass, field
from typing import List, Optional, Tuple
from .._base import percent_field
from platypush.schemas.dataclasses import percent_field
@dataclass

View file

@ -1,12 +1,11 @@
from marshmallow_dataclass import class_schema
from platypush.schemas.dataclasses import DataClassSchema
from .._base import SystemBaseSchema
from ._base import CpuInfoBaseSchema, CpuTimesBaseSchema
from ._model import CpuFrequency, CpuInfo, CpuStats, CpuTimes
CpuFrequencySchema = class_schema(CpuFrequency, base_schema=DataClassSchema)
CpuFrequencySchema = class_schema(CpuFrequency, base_schema=SystemBaseSchema)
CpuInfoSchema = class_schema(CpuInfo, base_schema=CpuInfoBaseSchema)
CpuTimesSchema = class_schema(CpuTimes, base_schema=CpuTimesBaseSchema)
CpuStatsSchema = class_schema(CpuStats, base_schema=DataClassSchema)
CpuStatsSchema = class_schema(CpuStats, base_schema=SystemBaseSchema)

View file

@ -1,15 +1,17 @@
from marshmallow import pre_load
from platypush.schemas.dataclasses import DataClassSchema
from .._base import SystemBaseSchema
class DiskBaseSchema(DataClassSchema):
class DiskBaseSchema(SystemBaseSchema):
"""
Base schema for disk stats.
"""
@pre_load
def pre_load(self, data: dict, **_) -> dict:
data = super().pre_load(data)
# Convert read/write/busy times from milliseconds to seconds
for attr in ['read_time', 'write_time', 'busy_time']:
if data.get(attr) is not None:

View file

@ -1,7 +1,7 @@
from dataclasses import dataclass, field
from typing import Optional
from .._base import percent_field
from platypush.schemas.dataclasses import percent_field
@dataclass

View file

@ -1,15 +1,17 @@
from marshmallow import pre_load
from platypush.schemas.dataclasses import DataClassSchema
from .._base import SystemBaseSchema
class MemoryStatsBaseSchema(DataClassSchema):
class MemoryStatsBaseSchema(SystemBaseSchema):
"""
Base schema for memory stats.
"""
@pre_load
def pre_load(self, data: dict, **_) -> dict:
data = super().pre_load(data)
# Normalize the percentage between 0 and 1
if data.get('percent') is not None:
data['percent'] /= 100

View file

@ -1,6 +1,6 @@
from dataclasses import dataclass, field
from .._base import percent_field
from platypush.schemas.dataclasses import percent_field
@dataclass

View file

@ -1,15 +1,16 @@
from marshmallow import pre_load
from platypush.schemas.dataclasses import DataClassSchema
from .._base import SystemBaseSchema
class NetworkInterfaceBaseSchema(DataClassSchema):
class NetworkInterfaceBaseSchema(SystemBaseSchema):
"""
Base schema for network interface stats.
"""
@pre_load
def pre_load(self, data: dict, **_) -> dict:
data = super().pre_load(data)
for in_attr, out_attr in {
'errin': 'errors_in',
'errout': 'errors_out',