pocketutils.tools.sys_tools

Complexity

Total Complexity

32

Size/Duplication

Total Lines	242
Duplicated Lines	0 %

Importance

Changes

0

13 Methods

Rating	Name	Duplication	Size	Complexity
A	SystemUtils.serialize_exception()	0	5	1
A	SystemUtils.is_linux()	0	2	1
A	Frame.as_dict()	0	2	1
A	SystemUtils.trace_exit()	0	5	1
A	ExitHandler.__call__()	0	5	2
A	SystemUtils.serialize_exception_msg()	0	3	1
A	SystemUtils.trace_signals()	0	8	2
D	SystemUtils.get_env_info()	0	90	12
B	SystemUtils.build_traceback()	0	23	5
A	SignalHandler.__call__()	0	5	2
A	SystemUtils.list_package_versions()	0	13	2
A	SystemUtils.is_macos()	0	2	1
A	SystemUtils.is_windows()	0	2	1

# SPDX-FileCopyrightText: Copyright 2020-2023, Contributors to pocketutils
# SPDX-PackageHomePage: https://github.com/dmyersturnbull/pocketutils
# SPDX-License-Identifier: Apache-2.0
"""
Low-level tools (e.g. memory management).
"""

import atexit
import importlib
import importlib.metadata
import locale
import logging
import os
import platform
import signal
import socket
import struct
import sys
import traceback
from collections.abc import Callable, Mapping, Sequence
from dataclasses import asdict, dataclass
from datetime import UTC, datetime
from getpass import getuser
from typing import Any, NamedTuple, Self

from pocketutils.core.input_output import Writeable

__all__ = ["Frame", "SerializedException", "SignalHandler", "ExitHandler", "SystemUtils", "SystemTools"]

logger = logging.getLogger("pocketutils")


@dataclass(frozen=True, slots=True, order=True, kw_only=True)
class Frame:
    depth: int
    filename: str
    line: int
    name: str
    repeats: int

    def as_dict(self: Self) -> Mapping[str, int | str]:
        return asdict(self)


class SerializedException(NamedTuple):
    message: Sequence[str]
    stacktrace: Sequence[Frame]


@dataclass(frozen=True, slots=True)
class SignalHandler:
    name: str
    code: int
    desc: str
    sink: Writeable | Callable[[str], Any]

    def __call__(self: Self) -> None:
        self.sink.write(f"~~{self.name}[{self.code}] ({self.desc})~~")
        traceback.print_stack(file=self.sink)
        for line in traceback.format_stack():
            self.sink.write(line)


@dataclass(frozen=True, slots=True)
class ExitHandler:
    sink: Writeable

    def __call__(self: Self) -> None:
        self.sink.write("~~EXIT~~")
        traceback.print_stack(file=self.sink)
        for line in traceback.format_stack():
            self.sink.write(line)


@dataclass(slots=True, frozen=True)
class SystemUtils:
    def is_linux(self: Self) -> bool:
        return sys.platform == "linux"

    def is_windows(self: Self) -> bool:
        return sys.platform == "win32"

    def is_macos(self: Self) -> bool:
        return sys.platform == "darwin"

    def get_env_info(self: Self, *, extended: bool = False, insecure: bool = False) -> Mapping[str, str]:
        """
        Get a dictionary of some system and environment information.
        Includes os_release, hostname, username, mem + disk, shell, etc.

        Args:
            extended: Get info from psutil
            insecure: Include data like hostname and username

        Warning:
            Even with `include_insecure=False`, avoid exposing this data to untrusted
            sources. For example, this includes the specific OS release, which could
            be used in attack.
        """

        now = datetime.now(UTC).astimezone().isoformat()
        uname = platform.uname()
        lang_code, encoding = locale.getlocale()
        # build up this dict:
        data = {}

        def _try(os_fn, k: str, *args):
            if any(a is None for a in args):
                return None
            try:
                v = os_fn(*args)
                data[k] = v
                return v
            except (OSError, ImportError):
                return None

        data.update(
            {
                "platform": platform.platform(),
                "python": ".".join(str(i) for i in sys.version_info),
                "os": uname.system,
                "os_release": uname.release,
                "os_version": uname.version,
                "machine": uname.machine,
                "byte_order": sys.byteorder,
                "processor": uname.processor,
                "build": sys.version,
                "python_bits": 8 * struct.calcsize("P"),
                "environment_info_capture_datetime": now,
                "encoding": encoding,
                "lang_code": lang_code,
                "recursion_limit": sys.getrecursionlimit(),
                "float_info": sys.float_info,
                "int_info": sys.int_info,
                "flags": sys.flags,
                "hash_info": sys.hash_info,
                "implementation": sys.implementation,
                "switch_interval": sys.getswitchinterval(),
                "filesystem_encoding": sys.getfilesystemencoding(),
            },
        )
        if "LANG" in os.environ:
            data["lang"] = os.environ["LANG"]
        if "SHELL" in os.environ:
            data["shell"] = os.environ["SHELL"]
        if "LC_ALL" in os.environ:
            data["lc_all"] = os.environ["LC_ALL"]
        if hasattr(sys, "winver"):
            data["win_ver"] = sys.getwindowsversion()
        if hasattr(sys, "mac_ver"):
            data["mac_ver"] = sys.mac_ver()
        if hasattr(sys, "linux_distribution"):
            data["linux_distribution"] = sys.linux_distribution()
        if insecure:
            _try(getuser, "username")
            _try(os.getlogin, "login")
            _try(socket.gethostname, "hostname")
            _try(os.getcwd, "cwd")
            pid = _try(os.getpid, "pid")
            ppid = _try(os.getppid, "parent_pid")
            if hasattr(os, "getpriority"):
                _try(os.getpriority, "priority", os.PRIO_PROCESS, pid)
                _try(os.getpriority, "parent_priority", os.PRIO_PROCESS, ppid)
        if extended:
            import psutil

            data.update(
                {
                    "disk_used": psutil.disk_usage(".").used,
                    "disk_free": psutil.disk_usage(".").free,
                    "memory_used": psutil.virtual_memory().used,
                    "memory_available": psutil.virtual_memory().available,
                },
            )
        return {k: str(v) for k, v in dict(data).items()}

    def list_package_versions(self: Self) -> Mapping[str, str]:
        """
        Returns installed packages and their version numbers.
        Reliable; uses importlib (Python 3.8+).
        """
        # calling .metadata reads the metadata file
        # and .version is an alias to .metadata["version"]
        # so make sure to only read once
        dct = {}
        for d in importlib.metadata.distributions():
            meta = d.metadata
            dct[meta["name"]] = meta["version"]
        return dct

    def serialize_exception(self: Self, e: BaseException) -> SerializedException:
        tbe = traceback.TracebackException.from_exception(e)
        msg = list(tbe.format_exception_only())
        tb = self.build_traceback(e)
        return SerializedException(msg, tb)

    def serialize_exception_msg(self: Self, e: BaseException) -> Sequence[str]:
        tbe = traceback.TracebackException.from_exception(e)
        return list(tbe.format_exception_only())

    def build_traceback(self: Self, e: BaseException) -> Sequence[Frame]:
        tb = []
        current = None
        tbe = traceback.TracebackException.from_exception(e)
        last, repeats = None, 0
        for i, s in enumerate(tbe.stack):
            current = Frame(depth=i, filename=s.filename, line=int(s.line), name=s.name, repeats=-1)
            if current == last:
                repeats += 1
            else:
                current = Frame(
                    depth=current.depth,
                    filename=current.filename,
                    line=current.line,
                    name=current.name,
                    repeats=repeats,
                )
                tb.append(current)
                repeats = 0
            last = current
        if current is not None and current == last:
            tb.append(current)
        return tb

    def trace_signals(self: Self, sink: Writeable = sys.stderr) -> None:
        """
        Registers signal handlers for all signals that log the traceback.
        Uses `signal.signal`.
        """
        for sig in signal.valid_signals():
            handler = SignalHandler(sig.name, sig.value, signal.strsignal(sig), sink)
            signal.signal(sig.value, handler)

    def trace_exit(self: Self, sink: Writeable = sys.stderr) -> None:
        """
        Registers an exit handler via `atexit.register` that logs the traceback.
        """
        atexit.register(ExitHandler(sink))


SystemTools = SystemUtils()