Skip to content

Commit

Permalink
Try using built-in methods to retrieve memory info
Browse files Browse the repository at this point in the history
  • Loading branch information
dormant-user committed Jan 1, 2025
1 parent 91a4c4d commit 950eb21
Show file tree
Hide file tree
Showing 4 changed files with 110 additions and 29 deletions.
7 changes: 3 additions & 4 deletions pyarchitecture/cpu/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,8 +2,6 @@
import os
from typing import Dict

import psutil

from pyarchitecture import models
from pyarchitecture.cpu import main

Expand Down Expand Up @@ -37,8 +35,9 @@ def get_cpu_info(cpu_lib: str | os.PathLike = None) -> Dict[str, int | str]:
Returns CPU name.
"""
cpu_name = main.get_name(_get_cpu_lib(cpu_lib))
cpu_count = os.cpu_count()
return {
"name": cpu_name,
"logical_cores": psutil.cpu_count(logical=True),
"physical_cores": psutil.cpu_count(logical=False),
"logical_cores": cpu_count,
"physical_cores": cpu_count / 2 if cpu_count >= 2 else 1,
}
112 changes: 94 additions & 18 deletions pyarchitecture/memory.py
Original file line number Diff line number Diff line change
@@ -1,25 +1,101 @@
import psutil
import logging
import os
import subprocess

from pyarchitecture import squire
from pyarchitecture import models

LOGGER = logging.getLogger(__name__)

def get_memory_info(raw: bool = False) -> dict[str, int | str]:
"""Get memory information for the host system.

Returns:
Dict[str, str]:
Returns memory information.
"""
if raw:
def get_memory_info_linux(mem_lib: str | os.PathLike):
memory_info = {}
with open(mem_lib) as f:
for line in f:
if line.startswith(('MemTotal', 'MemFree', 'MemAvailable', 'Buffers', 'Cached')):
parts = line.split()
memory_info[parts[0][:-1]] = int(parts[1]) # Convert the memory value to int (in kB)

# Convert values to bytes (kB to bytes)
total = memory_info.get('MemTotal', 0) * 1024
free = memory_info.get('MemFree', 0) * 1024
available = memory_info.get('MemAvailable', 0) * 1024
used = total - free - available

return {'total': total, 'free': free, 'used': used}


def get_memory_info_macos(mem_lib: str | os.PathLike):
def get_sysctl_value(key):
result = subprocess.run([mem_lib, key], capture_output=True, text=True)
if result.stdout.strip():
return int(result.stdout.split(":")[1].strip())
return 0

total = get_sysctl_value('hw.memsize')
free = get_sysctl_value('vm.page_free_count') * get_sysctl_value('hw.pagesize')
used = total - free

return {'total': total, 'free': free, 'used': used}


def get_memory_info_windows():
import ctypes
try:
memory_info = ctypes.windll.kernel32.GlobalMemoryStatusEx
memstatus = ctypes.create_string_buffer(64)
memory_info(ctypes.byref(memstatus))
total = ctypes.cast(memstatus[0], ctypes.c_ulonglong).value
free = ctypes.cast(memstatus[1], ctypes.c_ulonglong).value
"TypeError: cast() argument 2 must be a pointer type, not c_ulonglong"
used = total - free
return {'total': total, 'free': free, 'used': used}
except Exception:
class MEMORYSTATUSEX(ctypes.Structure):
_fields_ = [
("dwLength", ctypes.c_uint),
("dwMemoryLoad", ctypes.c_uint),
("ullTotalPhys", ctypes.c_ulonglong),
("ullAvailPhys", ctypes.c_ulonglong),
("ullTotalPageFile", ctypes.c_ulonglong),
("ullAvailPageFile", ctypes.c_ulonglong),
("ullTotalVirtual", ctypes.c_ulonglong),
("ullAvailVirtual", ctypes.c_ulonglong),
("sullAvailExtendedVirtual", ctypes.c_ulonglong)
]

# Initialize the MEMORYSTATUSEX structure
memory_status = MEMORYSTATUSEX()
memory_status.dwLength = ctypes.sizeof(MEMORYSTATUSEX)

# Call GlobalMemoryStatusEx to fill in the memory_status structure
if ctypes.windll.kernel32.GlobalMemoryStatusEx(ctypes.byref(memory_status)) == 0:
raise Exception("Failed to retrieve memory status")

# Extract the values from the structure
total = memory_status.ullTotalPhys # Total physical memory (in bytes)
available = memory_status.ullAvailPhys # Available physical memory (in bytes)
used = total - available # Used memory (in bytes)

# Optionally, you can also include virtual memory information
virtual_total = memory_status.ullTotalVirtual
virtual_available = memory_status.ullAvailVirtual

return {
"total": psutil.virtual_memory().total,
"available": psutil.virtual_memory().available,
"used": psutil.virtual_memory().used,
"free": psutil.virtual_memory().free,
'total': total,
'available': available,
'used': used,
'virtual_total': virtual_total,
'virtual_available': virtual_available,
}
return {
"total": squire.size_converter(psutil.virtual_memory().total),
"available": squire.size_converter(psutil.virtual_memory().available),
"used": squire.size_converter(psutil.virtual_memory().used),
"free": squire.size_converter(psutil.virtual_memory().free),


def get_memory_info(mem_lib: str | os.PathLike):
os_map = {
models.OperatingSystem.darwin: get_memory_info_macos,
models.OperatingSystem.linux: get_memory_info_linux,
models.OperatingSystem.windows: get_memory_info_windows
}
try:
return os_map[models.OPERATING_SYSTEM](mem_lib)
except Exception as error:
LOGGER.error(error)
17 changes: 13 additions & 4 deletions pyarchitecture/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,10 +5,10 @@
except ImportError:
from enum import Enum


class StrEnum(str, Enum):
"""Custom StrEnum object for python3.10."""


OPERATING_SYSTEM = platform.system().lower()


Expand All @@ -25,16 +25,25 @@ class OperatingSystem(StrEnum):


if OPERATING_SYSTEM not in (
OperatingSystem.linux,
OperatingSystem.darwin,
OperatingSystem.windows,
OperatingSystem.linux,
OperatingSystem.darwin,
OperatingSystem.windows,
):
raise RuntimeError(
f"{OPERATING_SYSTEM!r} is unsupported.\n\t"
"Host machine should either be macOS, Windows or any Linux distros"
)


def default_mem_lib():
"""Returns the default memory library dedicated to linux and macOS."""
return dict(
linux="/proc/meminfo",
darwin="/usr/sbin/sysctl",
windows="" # placeholder
)


def default_disk_lib():
"""Returns the default disks' library dedicated to each supported operating system."""
return dict(
Expand Down
3 changes: 0 additions & 3 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -14,9 +14,6 @@ classifiers = [
]
keywords = ["physical-drives", "PyArchitecture"]
requires-python = ">=3.10"
dependencies = [
"psutil>=6.1.1"
]

[tool.setuptools]
packages = [
Expand Down

0 comments on commit 950eb21

Please sign in to comment.