updated get_assignments

reidoko · Nov 14, 2024 · a2582c1 · a2582c1
1 parent 5fddd63
commit a2582c1
Showing 1 changed file with 13 additions and 21 deletions.
diff --git a/src/gradescope_api/course.py b/src/gradescope_api/course.py
@@ -71,30 +71,22 @@ def get_assignment(
         assignment_id = assignment_id or get_url_id(url=assignment_url, kind="assignments")
         return GradescopeAssignment(_client=self._client, _course=self, assignment_id=assignment_id)
 
-    def get_all_assignments(self) -> list[GradescopeAssignment]:
+    def get_assignments(self, where=lambda x: True) -> list[GradescopeAssignment]:
         response = self._client.session.get(f"https://www.gradescope.com/courses/{self.course_id}/assignments")
         soup = BeautifulSoup(response.content, "html.parser")
         props = soup.find(
             attrs={"data-react-class" : "AssignmentsTable"}
         )["data-react-props"]
         assignment_data = json.loads(props)["table_data"]
-        return [
-            self.get_assignment(assignment_url=f"https://www.gradescope.com/courses/{data['url']}") 
-            for data in filter(lambda x: x["type"] == "assignment", assignment_data)
-        ]
-
-    def get_assignments_where(self, match_fn) -> list[GradescopeAssignment]:
-        """
-        Get a list of assignments where the titles satisfy the match_fn
-        """
-        response = self._client.session.get(f"https://www.gradescope.com/courses/{self.course_id}/assignments")
-        soup = BeautifulSoup(response.content, "html.parser")
-        props = soup.find(
-            attrs={"data-react-class" : "AssignmentsTable"}
-        )["data-react-props"]
-        assignment_data = json.loads(props)["table_data"]
-        return [
-            self.get_assignment(assignment_url=f"https://www.gradescope.com/courses/{data['url']}") 
-            for data in filter(lambda x: x["type"] == "assignment" and match_fn(x['title']), assignment_data)
-        ]
-
+
+        if (callable(where)):
+            return [
+                self.get_assignment(assignment_url=f"https://www.gradescope.com/courses/{data['url']}") 
+                for data in filter(lambda x: x["type"] == "assignment" and where(x['title']), assignment_data)
+            ]
+        else:
+            key = where.lower()
+            return [
+                self.get_assignment(assignment_url=f"https://www.gradescope.com/courses/{data['url']}") 
+                for data in filter(lambda x: x["type"] == "assignment" and key in x['title'].lower(), assignment_data)
+            ]