From b71e9a997e288c18b1a169349d7a4e9792d8f2f2 Mon Sep 17 00:00:00 2001
From: epwalsh <petew@allenai.org>
Date: Thu, 13 Jun 2024 13:25:46 -0700
Subject: [PATCH] Add `author` filter to `Beaker.job.list()`

And add `Beaker.experiment.list()` method.
---
 CHANGELOG.md                  |  5 +++++
 beaker/services/experiment.py | 39 +++++++++++++++++++++++++++++++++++
 beaker/services/job.py        | 21 +++++++++++++------
 3 files changed, 59 insertions(+), 6 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 21f5739..700666a 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -9,6 +9,11 @@ use patch releases for compatibility fixes instead.
 
 ## Unreleased
 
+### Added
+
+- Added `author` filter to `Beaker.job.list()`.
+- Added `Beaker.experiment.list()` method.
+
 ## [v1.28.0](https://github.com/allenai/beaker-py/releases/tag/v1.28.0) - 2024-06-13
 
 ### Added
diff --git a/beaker/services/experiment.py b/beaker/services/experiment.py
index bde5426..7db9c41 100644
--- a/beaker/services/experiment.py
+++ b/beaker/services/experiment.py
@@ -136,6 +136,45 @@ def _parse_create_args(
         assert spec is not None
         return spec, name, workspace
 
+    def list(
+        self,
+        *,
+        author: Optional[Union[str, Account]] = None,
+        cluster: Optional[Union[str, Cluster]] = None,
+        finalized: bool = False,
+        node: Optional[Union[str, Node]] = None,
+    ) -> List[Experiment]:
+        """
+        List experiments.
+
+        :param author: List only experiments by particular author.
+        :param cluster: List experiments on a particular cluster.
+        :param finalized: List only finalized or non-finalized experiments.
+        :param node: List experiments on a particular node.
+
+        .. important::
+            Either ``cluster``, ``author``, ``experiment``, or ``node`` must be specified.
+            If ``node`` is specified, neither ``cluster`` nor ``experiment`` can be
+            specified.
+
+        :raises ValueError: If the arguments are invalid, e.g. both ``node`` and
+            ``cluster`` are specified.
+        :raises AccountNotFound: If the specified author doesn't exist.
+        :raises ClusterNotFound: If the specified cluster doesn't exist.
+        :raises NodeNotFound: If the specified node doesn't exist.
+        """
+        jobs = self.beaker.job.list(
+            author=author, cluster=cluster, finalized=finalized, kind=JobKind.execution, node=node
+        )
+        experiments: List[Experiment] = []
+        exp_ids: Set[str] = set()
+        for job in jobs:
+            assert job.execution is not None
+            if (exp_id := job.execution.experiment) not in exp_ids:
+                exp_ids.add(exp_id)
+                experiments.append(self.get(exp_id))
+        return experiments
+
     def create(
         self,
         *args,
diff --git a/beaker/services/job.py b/beaker/services/job.py
index 2e930d6..5348400 100644
--- a/beaker/services/job.py
+++ b/beaker/services/job.py
@@ -36,6 +36,8 @@ def get(self, job_id: str) -> Job:
 
     def list(
         self,
+        *,
+        author: Optional[Union[str, Account]] = None,
         cluster: Optional[Union[str, Cluster]] = None,
         experiment: Optional[Union[str, Experiment]] = None,
         finalized: bool = False,
@@ -45,37 +47,44 @@ def list(
         """
         List jobs.
 
-        :param cluster: List jobs on a cluster.
+        :param author: List only jobs by particular author.
+        :param cluster: List jobs on a particular cluster.
         :param experiment: List jobs in an experiment.
         :param finalized: List only finalized or non-finalized jobs.
         :param kind: List jobs of a certain kind.
-        :param node: List jobs on a node.
+        :param node: List jobs on a particular node.
 
         .. important::
-            Either ``cluster``, ``experiment``, or ``node`` must be specified.
+            Either ``cluster``, ``author``, ``experiment``, or ``node`` must be specified.
             If ``node`` is specified, neither ``cluster`` nor ``experiment`` can be
             specified.
 
         :raises ValueError: If the arguments are invalid, e.g. both ``node`` and
             ``cluster`` are specified.
+        :raises AccountNotFound: If the specified author doesn't exist.
         :raises ClusterNotFound: If the specified cluster doesn't exist.
         :raises ExperimentNotFound: If the specified experiment doesn't exist.
         :raises NodeNotFound: If the specified node doesn't exist.
         """
+        if node is None and cluster is None and experiment is None and author is None:
+            raise ValueError("You must specify one of 'node', 'cluster', 'experiment', or 'author'")
+
         # Validate arguments.
         if node is not None:
             if cluster is not None:
                 raise ValueError("You cannot specify both 'node' and 'cluster'")
             if experiment is not None:
                 raise ValueError("You cannot specify both 'node' and 'experiment'")
-        else:
-            if cluster is None and experiment is None:
-                raise ValueError("You must specify one of 'node', 'experiment', or 'cluster'")
 
         jobs: List[Job] = []
 
         # Build request options.
         request_opts: Dict[str, Any] = {}
+        if author is not None:
+            author_id = (
+                author.id if isinstance(author, Account) else self.beaker.account.get(author).id
+            )
+            request_opts["author"] = author_id
         if cluster is not None:
             cluster_id = (
                 cluster.id if isinstance(cluster, Cluster) else self.beaker.cluster.get(cluster).id