Repository: troycomi/reportseff

Issue:
feat: add ability to set sacct's "--units" argument
In order to let users see the difference between their ReqMem and the MaxRSS, I'd like to set them both to the same unit. `sacct --units <unit>` seems to be an easy way to achieve that. Would that fit into reportseff? If yes I'd be happy to open a PR.

Discussion:

Patch:
diff --git a/.github/workflows/coverage.yml b/.github/workflows/coverage.yml
index 72d0629..6ca8164 100644
--- a/.github/workflows/coverage.yml
+++ b/.github/workflows/coverage.yml
@@ -1,5 +1,5 @@
 name: Coverage
-on: push
+on: [push, pull_request]
 jobs:
   coverage:
     runs-on: ubuntu-latest
diff --git a/pyproject.toml b/pyproject.toml
index 67edd59..973baa7 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "reportseff"
-version = "2.5.0"
+version = "2.6.0"
 description= "Tablular seff output"
 authors = ["Troy Comi <troycomi@gmail.com>"]
 license = "MIT"
diff --git a/src/reportseff/console.py b/src/reportseff/console.py
index 047bc3e..df8fae4 100644
--- a/src/reportseff/console.py
+++ b/src/reportseff/console.py
@@ -56,6 +56,11 @@ from .parameters import ReportseffParameters
     default="",
     help="Only include jobs with the specified partition",
 )
+@click.option(
+    "--extra-args",
+    default="",
+    help="Extra arguments to forward to sacct",
+)
 @click.option(
     "-s", "--state", default="", help="Only include jobs with the specified states"
 )
@@ -65,7 +70,7 @@ from .parameters import ReportseffParameters
 @click.option(
     "--since",
     default="",
-    help="Only include jobs before this time. Can be valid sacct "
+    help="Only include jobs after this time. Can be valid sacct "
     "or as a comma separated list of time deltas, e.g. d=2,h=1 "
     "means 2 days, 1 hour before current time. Weeks, days, "
     "hours, and minutes can use case-insensitive abbreviations. "
@@ -139,6 +144,9 @@ def get_jobs(args: ReportseffParameters) -> Tuple[str, int]:
     inquirer.set_until(args.until)
 
     inquirer.set_partition(args.partition)
+
+    inquirer.set_extra_args(args.extra_args)
+
     add_jobs = False
 
     try:
diff --git a/src/reportseff/db_inquirer.py b/src/reportseff/db_inquirer.py
index e3401ac..f7587d4 100644
--- a/src/reportseff/db_inquirer.py
+++ b/src/reportseff/db_inquirer.py
@@ -2,6 +2,7 @@
 from abc import ABC, abstractmethod
 import datetime
 import re
+import shlex
 import subprocess
 from typing import Callable, Dict, List, Optional, Set
 
@@ -73,6 +74,14 @@ class BaseInquirer(ABC):
             partition: partition name
         """
 
+    @abstractmethod
+    def set_extra_args(self, extra_args: str) -> None:
+        """Set extra arguments to be forwarded to sacct.
+
+        Args:
+            extra_args: list of arguments
+        """
+
     @abstractmethod
     def all_users(self) -> None:
         """Ignore provided jobs, query for all users."""
@@ -152,6 +161,7 @@ class SacctInquirer(BaseInquirer):
         self.until: Optional[str] = None
         self.query_all_users: bool = False
         self.partition: Optional[str] = None
+        self.extra_args: Optional[str] = None
 
     def get_valid_formats(self) -> List[str]:
         """Get the valid formatting options supported by the inquirer.
@@ -191,18 +201,20 @@ class SacctInquirer(BaseInquirer):
             if not self.since:
                 start_date = datetime.date.today() - datetime.timedelta(days=7)
                 self.since = start_date.strftime("%m%d%y")  # MMDDYY
-            args += [f"--user={self.user}", f"--starttime={self.since}"]
+            args += [f"--user={self.user}"]
         elif self.query_all_users:
-            args += ["--allusers", f"--starttime={self.since}"]
+            args += ["--allusers"]
         else:
             args += ["--jobs=" + ",".join(jobs)]
-            if self.since:
-                args += [f"--starttime={self.since}"]
 
+        if self.since:
+            args += [f"--starttime={self.since}"]
         if self.partition:
             args += [f"--partition={self.partition}"]
         if self.until:
             args += [f"--endtime={self.until}"]
+        if self.extra_args:
+            args += shlex.split(self.extra_args)
         return args
 
     def get_db_output(
@@ -274,6 +286,14 @@ class SacctInquirer(BaseInquirer):
         """
         self.partition = partition
 
+    def set_extra_args(self, extra_args: str) -> None:
+        """Set extra arguments to be forwarded to sacct.
+
+        Args:
+            extra_args: list of arguments
+        """
+        self.extra_args = extra_args
+
     def all_users(self) -> None:
         """Query for all users if `since` is set."""
         self.query_all_users = True
diff --git a/src/reportseff/parameters.py b/src/reportseff/parameters.py
index 44cb352..4be0f29 100644
--- a/src/reportseff/parameters.py
+++ b/src/reportseff/parameters.py
@@ -22,6 +22,7 @@ class ReportseffParameters:
     state: str = ""
     user: str = ""
     partition: str = ""
+    extra_args: str = ""
 
     def __init__(
         self,
@@ -39,6 +40,7 @@ class ReportseffParameters:
         state: str = "",
         user: str = "",
         partition: str = "",
+        extra_args: str = "",
     ) -> None:
         """Create a new parameter object.
 
@@ -57,6 +59,7 @@ class ReportseffParameters:
             state: stats to display
             user: the user to report for
             partition: the partition to report for
+            extra_args: extra arguments to forward to sacct
         """
         self.color = color
         self.debug = debug
@@ -72,6 +75,7 @@ class ReportseffParameters:
         self.state = state
         self.user = user
         self.partition = partition
+        self.extra_args = extra_args
 
         if self.format_str.startswith("+"):
             self.format_str = (


Test patch:
diff --git a/tests/test_db_inquirer.py b/tests/test_db_inquirer.py
index b37edd9..a204af0 100644
--- a/tests/test_db_inquirer.py
+++ b/tests/test_db_inquirer.py
@@ -713,3 +713,17 @@ def test_partition_timelimit_issue_11(sacct, mocker):
         "mainqueue": "UNLIMITED",
         "mainqueue2": "10-00:00:00",
     }
+
+
+def test_extra_args_setting(sacct):
+    """Setting extra args are properly handled."""
+    sacct.set_extra_args('-D --units M --nodelist "node1 node2"')
+    assert sacct.extra_args == '-D --units M --nodelist "node1 node2"'
+    assert sacct.set_sacct_args(["123"]) == [
+        "--jobs=123",
+        "-D",
+        "--units",
+        "M",
+        "--nodelist",
+        "node1 node2",
+    ]
diff --git a/tests/test_reportseff.py b/tests/test_reportseff.py
index c628bd6..bc064dc 100644
--- a/tests/test_reportseff.py
+++ b/tests/test_reportseff.py
@@ -1,4 +1,5 @@
 """Test cli usage."""
+import shlex
 import subprocess
 
 from click.testing import CliRunner
@@ -1011,3 +1012,16 @@ def test_energy_reporting(mocker, mock_inquirer):
         "27",
     ]
     assert len(output) == 5
+
+
+def test_extra_args(mocker, mock_inquirer):
+    """Can add extra arguments for sacct."""
+    mocker.patch("reportseff.console.which", return_value=True)
+    runner = CliRunner()
+    mocker.patch("reportseff.console.get_jobs", return_value=("Testing", 1))
+    result = runner.invoke(
+        console.main,
+        shlex.split("--no-color --extra-args='-D --units M --nodelist=node1,node2'"),
+    )
+
+    assert result.exit_code == 0