From ef761a9ac4711fd38dbf75a8fabe231086b12f3c Mon Sep 17 00:00:00 2001 From: Thomas Ilsche Date: Fri, 2 Jun 2023 01:28:29 +0200 Subject: [PATCH] refactor: typing --- metricq_source_sysinfo/main.py | 7 +++--- metricq_source_sysinfo/sysinfo_source.py | 30 +++++++++++++++++------- 2 files changed, 24 insertions(+), 13 deletions(-) diff --git a/metricq_source_sysinfo/main.py b/metricq_source_sysinfo/main.py index 14d4aae..6ca338c 100644 --- a/metricq_source_sysinfo/main.py +++ b/metricq_source_sysinfo/main.py @@ -1,8 +1,7 @@ import logging import click - -import click_log +import click_log # type: ignore from metricq.logging import get_logger from .sysinfo_source import SysinfoSource @@ -19,7 +18,7 @@ @click.command() @click.argument("management-url", default="amqp://localhost/") @click.option("--token", default="source-sysinfo") -@click_log.simple_verbosity_option(logger) -def run(management_url, token): +@click_log.simple_verbosity_option(logger) # type: ignore +def run(management_url: str, token: str) -> None: src = SysinfoSource(management_url=management_url, token=token) src.run() diff --git a/metricq_source_sysinfo/sysinfo_source.py b/metricq_source_sysinfo/sysinfo_source.py index 72b06f2..d7eed6c 100644 --- a/metricq_source_sysinfo/sysinfo_source.py +++ b/metricq_source_sysinfo/sysinfo_source.py @@ -1,7 +1,9 @@ import asyncio -import socket import re +import socket +from typing import TYPE_CHECKING, Any, Optional +import metricq import psutil from metricq import IntervalSource, Timedelta, Timestamp, logging, rpc_handler @@ -12,15 +14,21 @@ # Ignore internal nics, especially from docker _NIC_IGNORE_PATTERN = re.compile(r"^(lo|br|docker|veth)") +if TYPE_CHECKING: + # Unfortunately there's no public type for it + from psutil._common import sdiskio, snetio + + class SysinfoSource(IntervalSource): - def __init__(self, *args, **kwargs): + prev_timestamp: Optional[metricq.Timestamp] = None + prev_net_io: Optional[dict[str, snetio]] = None + prev_disk_io: Optional[dict[str, sdiskio]] = None + + def __init__(self, *args: Any, **kwargs: Any) -> None: super().__init__(*args, client_version=client_version, **kwargs) - self.prev_timestamp = None - self.prev_net_io = None - self.prev_disk_io = None @rpc_handler("config") - async def _on_config(self, **config): + async def _on_config(self, **config: Any) -> None: logger.info("config: {}", config) rate = config["rate"] self.period = Timedelta.from_s(1 / rate) @@ -77,7 +85,7 @@ async def _on_config(self, **config): # Disk self.prev_disk_io = psutil.disk_io_counters(perdisk=True, nowrap=True) - for disk_name in self.prev_disk_io.keys(): + for disk_name in self.prev_disk_io.keys(): for rw in "read", "written": meta[f"disk.{disk_name}.{rw}.count"] = { "rate": rate, @@ -94,10 +102,12 @@ async def _on_config(self, **config): {self.prefix + key: value for key, value in meta.items()} ) - async def send(self, metric, timestamp, value): + async def send( + self, metric: str, timestamp: metricq.Timestamp, value: float + ) -> None: await super().send(self.prefix + metric, timestamp, value) - async def update(self): + async def update(self) -> None: assert self.prev_timestamp is not None, "update() called before _on_config()" now = Timestamp.now() @@ -114,6 +124,7 @@ async def update(self): net_io = psutil.net_io_counters(pernic=True, nowrap=True) duration_s = (now - self.prev_timestamp).s + assert self.prev_net_io is not None for nic_name, net_values in net_io.items(): if _NIC_IGNORE_PATTERN.match(nic_name): continue @@ -148,6 +159,7 @@ async def update(self): ) disk_io = psutil.disk_io_counters(perdisk=True, nowrap=True) + assert self.prev_disk_io is not None duration_s = (now - self.prev_timestamp).s for disk_name, disk_values in disk_io.items(): prev_disk_values = self.prev_disk_io[disk_name]