From 875372de6e1c4c1f63e3a27c62e3cf12f41a0585 Mon Sep 17 00:00:00 2001
From: vsivanandharao_expedia <vsivanandharao@expediagroup.com>
Date: Tue, 31 Dec 2024 23:03:15 -0600
Subject: [PATCH] Include memory information

---
 pyarchitecture/__init__.py           |  6 +++---
 pyarchitecture/cpu/__init__.py       |  7 ++++---
 pyarchitecture/disks/macOS.py        |  2 +-
 pyarchitecture/disks/windows.py      |  2 +-
 pyarchitecture/memory.py             | 25 +++++++++++++++++++++++++
 pyarchitecture/{disks => }/squire.py |  0
 pyproject.toml                       |  3 +++
 7 files changed, 37 insertions(+), 8 deletions(-)
 create mode 100644 pyarchitecture/memory.py
 rename pyarchitecture/{disks => }/squire.py (100%)

diff --git a/pyarchitecture/__init__.py b/pyarchitecture/__init__.py
index 6fb1ed3..15c6560 100644
--- a/pyarchitecture/__init__.py
+++ b/pyarchitecture/__init__.py
@@ -4,7 +4,7 @@
 import time
 from typing import Any, Dict
 
-from pyarchitecture import cpu, disks, gpu
+from pyarchitecture import cpu, disks, gpu, memory
 
 version = "0.0.0-a0"
 
@@ -20,6 +20,7 @@ def all_components() -> Dict[str, Any]:
         "Disks": disks.get_all_disks(),
         "CPU": cpu.get_cpu_info(),
         "GPU": gpu.get_gpu_names(),
+        "Memory": memory.get_memory_info(),
     }
 
 
@@ -107,10 +108,9 @@ def commandline() -> None:
                 data["GPU"] = gpu.get_gpu_names()
             if disk_info:
                 data["Disks"] = disks.get_all_disks()
-        import io
 
         with open(filename, "w") as json_file:
-            json.dump(data, json_file, indent=2)  # type: io.TextIOBase
+            json.dump(data, json_file, indent=2)
         print(f"Architecture information has been stored in {filename!r}")
         sys.exit(0)
     else:
diff --git a/pyarchitecture/cpu/__init__.py b/pyarchitecture/cpu/__init__.py
index 6ac9791..f43e25a 100644
--- a/pyarchitecture/cpu/__init__.py
+++ b/pyarchitecture/cpu/__init__.py
@@ -2,6 +2,8 @@
 import os
 from typing import Dict
 
+import psutil
+
 from pyarchitecture import models
 from pyarchitecture.cpu import main
 
@@ -35,9 +37,8 @@ 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": cpu_count,
-        "physical_cores": int(cpu_count / 2) if cpu_count >= 2 else 1,
+        "logical_cores": psutil.cpu_count(logical=True),
+        "physical_cores": psutil.cpu_count(logical=False),
     }
diff --git a/pyarchitecture/disks/macOS.py b/pyarchitecture/disks/macOS.py
index b9e1899..048ff3b 100644
--- a/pyarchitecture/disks/macOS.py
+++ b/pyarchitecture/disks/macOS.py
@@ -5,7 +5,7 @@
 from collections import defaultdict
 from typing import Dict, List
 
-from pyarchitecture.disks import squire
+from pyarchitecture import squire
 
 LOGGER = logging.getLogger(__name__)
 
diff --git a/pyarchitecture/disks/windows.py b/pyarchitecture/disks/windows.py
index 1d752da..86b9d17 100644
--- a/pyarchitecture/disks/windows.py
+++ b/pyarchitecture/disks/windows.py
@@ -6,7 +6,7 @@
 import subprocess
 from typing import Dict, List, Tuple
 
-from pyarchitecture.disks import squire
+from pyarchitecture import squire
 
 LOGGER = logging.getLogger(__name__)
 
diff --git a/pyarchitecture/memory.py b/pyarchitecture/memory.py
new file mode 100644
index 0000000..0aafe55
--- /dev/null
+++ b/pyarchitecture/memory.py
@@ -0,0 +1,25 @@
+import psutil
+
+from pyarchitecture import squire
+
+
+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:
+        return {
+            "total": psutil.virtual_memory().total,
+            "available": psutil.virtual_memory().available,
+            "used": psutil.virtual_memory().used,
+            "free": psutil.virtual_memory().free,
+        }
+    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),
+    }
diff --git a/pyarchitecture/disks/squire.py b/pyarchitecture/squire.py
similarity index 100%
rename from pyarchitecture/disks/squire.py
rename to pyarchitecture/squire.py
diff --git a/pyproject.toml b/pyproject.toml
index e7762c7..28ffd3c 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -14,6 +14,9 @@ classifiers = [
 ]
 keywords = ["physical-drives", "PyArchitecture"]
 requires-python = ">=3.10"
+dependencies = [
+    "psutil>=6.1.1"
+]
 
 [tool.setuptools]
 packages = [