Restructure project to avoid passing around env_config

README.md

@@ -36,7 +36,8 @@ pip install PyUdisk
 import pyudisk
 
 if __name__ == '__main__':
-    pyudisk.monitor()
+    for metric in pyudisk.smart_metrics():
+        print(metric)
 ```
 
 **CLI**

pyudisk/__init__.py

@@ -5,7 +5,7 @@
 
 import click
 
-from .main import EnvConfig, generate_report, monitor, smart_metrics  # noqa: F401
+from .main import generate_report, monitor, smart_metrics  # noqa: F401
 
 version = "1.1.0"
 

pyudisk/config.py

@@ -7,7 +7,7 @@
 from pydantic import BaseModel, DirectoryPath, Field, FilePath, HttpUrl, field_validator
 from pydantic_settings import BaseSettings
 
-from .models import udisk
+from . import models
 
 OPERATING_SYSTEM = platform.system()
 SMARTCTL_LIB = shutil.which("smartctl")
@@ -150,7 +150,7 @@ def validate_metrics(cls, value: Metric | List[Metric]) -> List[Metric]:
                 for dtype in dtypes.get("anyOf")
                 if dtype.get("type", "") != "null"
             ]
-            for name, dtypes in udisk.Attributes.model_json_schema()
+            for name, dtypes in models.udisk.Attributes.model_json_schema()
             .get("properties")
             .items()
         }
@@ -207,3 +207,6 @@ class Config:
 
         env_file = os.environ.get("env_file") or os.environ.get("ENV_FILE") or ".env"
         extra = "ignore"
+
+
+env: EnvConfig = EnvConfig  # noqa: TypeCheck

pyudisk/disk_data.py

@@ -5,10 +5,8 @@
 import pyarchitecture
 from psutil._common import sdiskpart
 
-from .config import EnvConfig
+from . import config, models, util
 from .logger import LOGGER
-from .models import SystemPartitions
-from .util import size_converter
 
 
 def get_partitions() -> Generator[sdiskpart]:
@@ -18,7 +16,7 @@ def get_partitions() -> Generator[sdiskpart]:
         sdiskpart:
         Yields the partition datastructure.
     """
-    system_partitions = SystemPartitions()
+    system_partitions = models.SystemPartitions()
     for partition in psutil.disk_partitions():
         is_not_system_mount = all(
             not partition.mountpoint.startswith(mnt)
@@ -53,7 +51,7 @@ def get_disk_io() -> Generator[Dict[str, str | List[str]]]:
         # If there is only one physical disk, then set the mountpoint to root (/)
         if len(physical_disks) == 1:
             yield dict(
-                size=size_converter(psutil.disk_usage("/").total),
+                size=util.size_converter(psutil.disk_usage("/").total),
                 device_id=physical_disks[0],
                 node=f"/dev/{physical_disks[0]}",
                 mountpoints=["/"],
@@ -62,7 +60,7 @@ def get_disk_io() -> Generator[Dict[str, str | List[str]]]:
             # If there are multiple physical disks, then set the mountpoint to disk path itself
             for physical_disk in physical_disks:
                 yield dict(
-                    size=size_converter(psutil.disk_usage("/").total),
+                    size=util.size_converter(psutil.disk_usage("/").total),
                     device_id=physical_disk,
                     node=f"/dev/{physical_disk}",
                     mountpoints=[f"/dev/{physical_disk}"],
@@ -78,18 +76,17 @@ def partitions() -> Generator[Dict[str, str | List[str]]]:
     """
     for partition in get_partitions():
         yield dict(
-            size=size_converter(psutil.disk_usage(partition.mountpoint).total),
+            size=util.size_converter(psutil.disk_usage(partition.mountpoint).total),
             device_id=partition.device.lstrip("/dev/"),
             node=partition.device,
             mountpoints=[partition.device],
         )
 
 
-def get_disk_data(env: EnvConfig, posix: bool) -> List[Dict[str, str | List[str]]]:
+def get_disk_data(posix: bool) -> List[Dict[str, str | List[str]]]:
     """Get disk information for macOS and Windows machines.
 
     Args:
-        env: Environment variables configuration.
         posix: If the operating system is POSIX compliant.
 
     Returns:
@@ -98,7 +95,7 @@ def get_disk_data(env: EnvConfig, posix: bool) -> List[Dict[str, str | List[str]
     """
     if posix:
         # 1: Attempt to extract physical disks from PyArchitecture
-        if pyarch := pyarchitecture.disks.get_all_disks(env.disk_lib):
+        if pyarch := pyarchitecture.disks.get_all_disks(config.env.disk_lib):
             return pyarch
         # 2: Assume disks with non-zero write count as physical disks
         # disk_io_counters fn will fetch disks rather than partitions (similar to output from 'diskutil list')

pyudisk/main.py

@@ -6,39 +6,34 @@
 
 from pydantic import NewPath, ValidationError
 
-from .config import OPERATING_SYSTEM, EnvConfig, OperationSystem
-from .disk_data import get_disk_data
+from . import config, disk_data, metrics, models, notification, parsers, util
 from .logger import LOGGER
-from .metrics import get_smart_metrics, get_udisk_metrics
-from .models import smartctl, udisk
-from .notification import notification_service, send_report
-from .parsers import parse_block_devices, parse_drives
-from .util import standard
 
 
-def smart_metrics(env: EnvConfig) -> Generator[udisk.Disk | smartctl.Disk]:
+def smart_metrics(**kwargs) -> Generator[models.udisk.Disk | models.smartctl.Disk]:
     """Gathers smart metrics using udisksctl dump, and constructs a Disk object.
 
-    Args:
-        env: Environment variables configuration.
-
     Yields:
         Disk:
         Yields the Disk object from the generated Dataframe.
     """
-    if OPERATING_SYSTEM in (OperationSystem.darwin, OperationSystem.windows):
-        for device_info in get_disk_data(
-            env, OPERATING_SYSTEM != OperationSystem.windows
+    config.env = config.EnvConfig(**kwargs)
+    if config.OPERATING_SYSTEM in (
+        config.OperationSystem.darwin,
+        config.OperationSystem.windows,
+    ):
+        for device_info in disk_data.get_disk_data(
+            config.OPERATING_SYSTEM != config.OperationSystem.windows
         ):
-            if metrics := get_smart_metrics(env.smart_lib, device_info):
+            if retrieved_metrics := metrics.get_smart_metrics(device_info):
                 try:
-                    yield smartctl.Disk(**metrics)
+                    yield models.smartctl.Disk(**retrieved_metrics)
                 except ValidationError as error:
                     LOGGER.error(error.errors())
         return
-    smart_dump = get_udisk_metrics(env)
-    block_devices = parse_block_devices(smart_dump)
-    drives = {k: v for k, v in sorted(parse_drives(smart_dump).items())}
+    smart_dump = metrics.get_udisk_metrics()
+    block_devices = parsers.parse_block_devices(smart_dump)
+    drives = {k: v for k, v in sorted(parsers.parse_drives(smart_dump).items())}
     diff = set()
     # Enable mount warning by default (log warning messages if disk is not mounted)
     mount_warning = os.environ.get("MOUNT_WARNING", "1") == "1"
@@ -57,18 +52,18 @@ def smart_metrics(env: EnvConfig) -> Generator[udisk.Disk | smartctl.Disk]:
             LOGGER.warning("UNmounted drive(s) found - '%s'", ", ".join(diff))
     optional_fields = [
         k
-        for k, v in udisk.Disk.model_json_schema().get("properties").items()
+        for k, v in models.udisk.Disk.model_json_schema().get("properties").items()
         if v.get("anyOf", [{}])[-1].get("type", "") == "null"
     ]
-    # S.M.A.R.T metrics can be null, but the keys are mandatory
+    # UDisk metrics can be null, but the keys are mandatory
     for drive in drives.values():
         for key in optional_fields:
             if key not in drive.keys():
                 drive[key] = None
     for drive, data in drives.items():
         if block_data := block_devices.get(drive):
             data["Partition"] = block_data
-            yield udisk.Disk(
+            yield models.udisk.Disk(
                 id=drive, model=data.get("Info", {}).get("Model", ""), **data
             )
         elif drive not in diff:
@@ -92,11 +87,11 @@ def generate_html(
     try:
         import jinja2
     except ModuleNotFoundError:
-        standard()
+        util.standard()
 
     template_dir = os.path.join(pathlib.Path(__file__).parent, "templates")
-    env = jinja2.Environment(loader=jinja2.FileSystemLoader(template_dir))
-    template = env.get_template(f"{OPERATING_SYSTEM}.html")
+    jinja_env = jinja2.Environment(loader=jinja2.FileSystemLoader(template_dir))
+    template = jinja_env.get_template(f"{config.OPERATING_SYSTEM}.html")
     now = datetime.now()
     html_output = template.render(
         data=data, last_updated=f"{now.strftime('%c')} {now.astimezone().tzinfo}"
@@ -118,9 +113,9 @@ def generate_report(**kwargs) -> str:
         str:
         Returns the report filepath.
     """
-    env = EnvConfig(**kwargs)
+    config.env = config.EnvConfig(**kwargs)
     if kwargs.get("raw"):
-        return generate_html([disk.model_dump() for disk in smart_metrics(env)])
+        return generate_html([disk.model_dump() for disk in smart_metrics()])
     if report_file := kwargs.get("filepath"):
         assert report_file.endswith(
             ".html"
@@ -129,35 +124,33 @@ def generate_report(**kwargs) -> str:
         os.makedirs(report_dir, exist_ok=True)
     else:
         if directory := kwargs.get("directory"):
-            env.report_dir = directory
-        os.makedirs(env.report_dir, exist_ok=True)
+            config.env.report_dir = directory
+        os.makedirs(config.env.report_dir, exist_ok=True)
         report_file = datetime.now().strftime(
-            os.path.join(env.report_dir, env.report_file)
+            os.path.join(config.env.report_dir, config.env.report_file)
         )
     LOGGER.info("Generating disk report")
-    disk_report = [disk.model_dump() for disk in smart_metrics(env)]
+    disk_report = [disk.model_dump() for disk in smart_metrics()]
     generate_html(disk_report, report_file)
     LOGGER.info("Report has been stored in %s", report_file)
     return report_file
 
 
-def monitor_disk(env: EnvConfig) -> Generator[udisk.Disk]:
+def monitor_disk(**kwargs) -> Generator[models.udisk.Disk]:
     """Monitors disk attributes based on the configuration.
 
-    Args:
-        env: Environment variables configuration.
-
     Yields:
         Disk:
         Data structure parsed as a Disk object.
     """
     assert (
-        OPERATING_SYSTEM == OperationSystem.linux
+        config.OPERATING_SYSTEM == config.OperationSystem.linux
     ), "Monitoring feature is available only for Linux machines!!"
+    config.env = config.EnvConfig(**kwargs)
     message = ""
-    for disk in smart_metrics(env):
+    for disk in smart_metrics():
         if disk.Attributes:
-            for metric in env.metrics:
+            for metric in config.env.metrics:
                 attribute = disk.Attributes.model_dump().get(metric.attribute)
                 if metric.max_threshold and attribute >= metric.max_threshold:
                     msg = f"{metric.attribute!r} for {disk.id!r} is >= {metric.max_threshold} at {attribute}"
@@ -175,9 +168,7 @@ def monitor_disk(env: EnvConfig) -> Generator[udisk.Disk]:
             LOGGER.warning("No attributes were loaded for %s", disk.model)
         yield disk
     if message:
-        notification_service(
-            title="Disk Monitor Alert!!", message=message, env_config=env
-        )
+        notification.notification_service(title="Disk Monitor Alert!!", message=message)
 
 
 def monitor(**kwargs) -> None:
@@ -187,27 +178,24 @@ def monitor(**kwargs) -> None:
         **kwargs: Arbitrary keyword arguments.
     """
     assert (
-        OPERATING_SYSTEM == OperationSystem.linux
+        config.OPERATING_SYSTEM == config.OperationSystem.linux
     ), "Monitoring feature is available only for Linux machines!!"
-    env = EnvConfig(**kwargs)
-    disk_report = [disk.model_dump() for disk in monitor_disk(env)]
+    config.env = config.EnvConfig(**kwargs)
+    disk_report = [disk.model_dump() for disk in monitor_disk()]
     if disk_report:
         LOGGER.info(
             "Disk monitor report has been generated for %d disks", len(disk_report)
         )
-        if env.disk_report:
-            os.makedirs(env.report_dir, exist_ok=True)
+        if config.env.disk_report:
+            os.makedirs(config.env.report_dir, exist_ok=True)
             report_file = datetime.now().strftime(
-                os.path.join(env.report_dir, env.report_file)
+                os.path.join(config.env.report_dir, config.env.report_file)
             )
             report_data = generate_html(disk_report, report_file)
-            if env.gmail_user and env.gmail_pass and env.recipient:
-                LOGGER.info("Sending an email disk report to %s", env.recipient)
-                send_report(
+            if config.env.gmail_user and config.env.gmail_pass and config.env.recipient:
+                LOGGER.info("Sending an email disk report to %s", config.env.recipient)
+                notification.send_report(
                     title=f"Disk Report - {datetime.now().strftime('%c')}",
-                    user=env.gmail_user,
-                    password=env.gmail_pass,
-                    recipient=env.recipient,
                     content=report_data,
                 )
             else:

pyudisk/metrics.py

@@ -2,20 +2,14 @@
 import subprocess
 from typing import Any, Dict, List
 
-from pydantic import FilePath
-
-from .config import EnvConfig
+from . import config, models
 from .logger import LOGGER
-from .models import smartctl
 
 
-def get_smart_metrics(
-    smart_lib: FilePath, device_info: Dict[str, str | List[str]]
-) -> Dict[str, Any]:
+def get_smart_metrics(device_info: Dict[str, str | List[str]]) -> Dict[str, Any]:
     """Gathers disk information using the 'smartctl' command.
 
     Args:
-        smart_lib: Library path to 'smartctl' command.
         device_info: Device information retrieved.
 
     Returns:
@@ -26,7 +20,7 @@ def get_smart_metrics(
     mountpoints = device_info["mountpoints"]
     try:
         result = subprocess.run(
-            [smart_lib, "-a", device_id, "--json"],
+            [config.env.smart_lib, "-a", device_id, "--json"],
             capture_output=True,
             text=True,
             check=False,
@@ -46,25 +40,23 @@ def get_smart_metrics(
         LOGGER.error("[%d]: %s", error.returncode, result)
         output = {}
     output["device"] = output.get(
-        "device", smartctl.Device(name=device_id, info_name=device_id).model_dump()
+        "device",
+        models.smartctl.Device(name=device_id, info_name=device_id).model_dump(),
     )
     output["model_name"] = output.get("model_name", device_info.get("name"))
     output["mountpoints"] = mountpoints
     return output
 
 
-def get_udisk_metrics(env: EnvConfig) -> str:
+def get_udisk_metrics() -> str:
     """Gathers disk information using the dump from 'udisksctl' command.
 
-    Args:
-        env: Environment variables configuration.
-
     Returns:
         str:
         Returns the output from disk util dump.
     """
     try:
-        output = subprocess.check_output(f"{env.smart_lib} dump", shell=True)
+        output = subprocess.check_output(f"{config.env.smart_lib} dump", shell=True)
     except subprocess.CalledProcessError as error:
         result = error.output.decode(encoding="UTF-8").strip()
         LOGGER.error(f"[{error.returncode}]: {result}\n")

pyudisk/models/__init__.py

@@ -2,6 +2,8 @@
 
 from pydantic import BaseModel, Field
 
+from . import smartctl, udisk  # noqa: F401
+
 
 class SystemPartitions(BaseModel):
     """System partitions' mountpoints and fstypes."""