Merge pull request #15 from mbdevpl/feature/improve-swap-support

improve swap support

Author: mbdevpl

README.rst

@@ -73,8 +73,8 @@ Using
 Installing *system-query* doesn't enable all the features by default. Some of the query functions
 will work only on **some** systems. To attempt installation with all features enables,
 run :bash:`pip3 install system-query[all]`. If something brakes, you can narrow down the features
-by typing a feature scope instead of :bash:`all`.
-You can choose from :bash:`cpu`, :bash:`gpu`, :bash:`hdd`, :bash:`ram` and :bash:`swap`.
+by typing a feature scope instead of ``all``.
+You can choose from ``cpu``, ``gpu``, ``hdd``, ``ram`` and ``swap``.
 E.g. :bash:`pip3 install system-query[gpu]`. You can also select more than one feature
 at the same time, e.g. :bash:`pip3 install system-query[cpu,hdd,ram]`.
 
@@ -140,13 +140,18 @@ assuming their executables are in system path.
 system_query.query_swap()
 ~~~~~~~~~~~~~~~~~~~~~~~~~
 
-To be able to see amount of swap space, install Python package :bash:`psutil`.
+To be able to see amount of swap space, install Python package ``psutil``.
+
+.. code:: python
+
+    In [9]: system_query.query_swap()
+    Out[9]: {'total': 0}
 
 
 As command-line tool
 --------------------
 
-For example:
+Below will run :python:`system_query.query_all()` and output results to stdout:
 
 .. code:: bash
 
@@ -161,23 +166,23 @@ For example:
      'host': 'TestMachine',
      'os': 'Linux-4.4.0-109-generic-x86_64-with-debian-stretch-sid',
      'ram': {'total': 33701269504},
-     'swap': 0}
+     'swap': {'total': 0}}
 
-Usage information:
+Please use ``-h`` to see usage information:
 
 .. code::
 
     $ python3 -m system_query -h
-    usage: system_query [-h] [-s {all,cpu,gpu,ram}] [-f {raw,json}] [-t TARGET]
-                        [--version]
+    usage: system_query [-h] [-s {all,cpu,gpu,ram,swap}] [-f {raw,json}]
+                        [-t TARGET] [--version]
 
     Comprehensive and concise system information tool. Query a given hardware
     and/or software scope of your system and get results in human- and machine-
     readable formats.
 
     optional arguments:
       -h, --help            show this help message and exit
-      -s {all,cpu,gpu,ram}, --scope {all,cpu,gpu,ram}
+      -s {all,cpu,gpu,ram,swap}, --scope {all,cpu,gpu,ram,swap}
                             Scope of the query (default: all)
       -f {raw,json}, --format {raw,json}
                             Format of the results of the query. (default: raw)
@@ -187,6 +192,8 @@ Usage information:
                             stdout and stderr, respectively. (default: stdout)
       --version             show program's version number and exit
 
+    Copyright 2017-2023 by the contributors, Apache License 2.0,
+    https://github.com/mbdevpl/system-query
 
 Requirements
 ============

system_query/__init__.py

@@ -2,7 +2,7 @@
 
 __all__ = [
     'query_all', 'query_cpu', 'query_gpus', 'query_hdds', 'query_ram', 'query_software',
-    'query_and_export']
+    'query_swap', 'query_and_export']
 
 from .all_info import query_all
 from .cpu_info import query_cpu
@@ -12,5 +12,5 @@
 # from .os_info import query_os
 from .ram_info import query_ram
 from .software_info import query_software
-# from .swap_info import query_swap
+from .swap_info import query_swap
 from .query import query_and_export

system_query/main.py

@@ -5,7 +5,7 @@
 import sys
 
 from ._version import VERSION
-from .query import query_and_export
+from .query import QUERY_FUNCTIONS, query_and_export
 
 
 def main(args=None, namespace=None):
@@ -21,7 +21,7 @@ def main(args=None, namespace=None):
         https://github.com/mbdevpl/system-query''',
         formatter_class=argparse.ArgumentDefaultsHelpFormatter, allow_abbrev=True)
     parser.add_argument(
-        '-s', '--scope', type=str, default='all', choices=['all', 'cpu', 'gpu', 'ram'],
+        '-s', '--scope', type=str, default='all', choices=list(QUERY_FUNCTIONS.keys()),
         help='''Scope of the query''')
     parser.add_argument(
         '-f', '--format', type=str, default='raw', choices=['raw', 'json'],

system_query/query.py

@@ -10,6 +10,7 @@
 from .cpu_info import query_cpu
 from .gpu_info import query_gpus
 from .ram_info import query_ram
+from .swap_info import query_swap
 
 JSON_INDENT = 2
 
@@ -21,28 +22,30 @@ def query_and_export(query_scope: str, export_format: str, export_target: t.Any,
 
     Currently implemented values are:
 
-    - query_scope: all, cpu, gpu, ram.
+    - query_scope: all, cpu, gpu, ram, swap.
     - export_format: json, raw.
     - export_target: sys.stdout, sys.stderr, path.
     """
     info = query(query_scope, **kwargs)
     export(info, export_format, export_target)
 
 
+QUERY_FUNCTIONS = {
+    'all': query_all,
+    'cpu': query_cpu,
+    'gpu': query_gpus,
+    'ram': query_ram,
+    'swap': query_swap,
+}
+
+
 def query(query_scope: str, **kwargs) -> t.Any:
     """Wrap around selected system query functions."""
-    info: t.Any
-    if query_scope == 'all':
-        info = query_all(**kwargs)
-    elif query_scope == 'cpu':
-        info = query_cpu(**kwargs)
-    elif query_scope == 'gpu':
-        info = query_gpus(**kwargs)
-    elif query_scope == 'ram':
-        info = query_ram(**kwargs)
-    else:
-        raise NotImplementedError(f'scope={query_scope}')
-    return info
+    try:
+        query_function = QUERY_FUNCTIONS[query_scope]
+    except KeyError as err:
+        raise NotImplementedError(f'scope={query_scope}') from err
+    return query_function(**kwargs)  # type: ignore
 
 
 def export(info, export_format: str, export_target: t.Any):

system_query/swap_info.py

@@ -5,9 +5,9 @@
 from .available_features import psutil, SWAP
 
 
-def query_swap() -> t.Optional[int]:
+def query_swap(**_) -> t.Optional[t.Dict[str, t.Any]]:
     """Get information about swap."""
     if not SWAP:
         return None
     total_swap = psutil.swap_memory().total
-    return total_swap
+    return {'total': total_swap}

test/test_available_features.py

@@ -2,8 +2,8 @@
 
 import unittest
 
-from system_query.available_features import CPU, GPU, HDD
-from system_query import query_cpu, query_gpus, query_hdds
+from system_query.available_features import CPU, GPU, HDD, SWAP
+from system_query import query_cpu, query_gpus, query_hdds, query_swap
 
 
 class Tests(unittest.TestCase):
@@ -19,3 +19,7 @@ def test_gpus(self):
     @unittest.skipUnless(HDD, 'querying HDDs is not supported')
     def test_hdds(self):
         self.assertGreaterEqual(len(query_hdds()), 1)
+
+    @unittest.skipUnless(SWAP, 'querying swap is not supported')
+    def test_swap(self):
+        self.assertGreaterEqual(len(query_swap()), 1)

test/test_query.py

@@ -16,8 +16,6 @@ def test_unsupported_scope(self):
             query_and_export('host', 'raw', sys.stdout)
         with self.assertRaises(NotImplementedError):
             query_and_export('os', 'raw', sys.stdout)
-        with self.assertRaises(NotImplementedError):
-            query_and_export('swap', 'raw', sys.stdout)
 
     def test_unsupported_format(self):
         with self.assertRaises(NotImplementedError):