diff --git a/src/ds_caselaw_utils/courts.py b/src/ds_caselaw_utils/courts.py index 4493319..6816495 100644 --- a/src/ds_caselaw_utils/courts.py +++ b/src/ds_caselaw_utils/courts.py @@ -74,6 +74,18 @@ def get_selectable(self): courts.append(Court(court)) return courts + def get_selectable_groups(self): + groups = [] + for category in self._data: + courts = [ + Court(court) + for court in category.get("courts") + if court.get("selectable") + ] + if len(courts) > 0: + groups.append(CourtGroup(category.get("display_name"), courts)) + return groups + def get_listable_groups(self): groups = [] for category in self._data: diff --git a/src/ds_caselaw_utils/test_courts.py b/src/ds_caselaw_utils/test_courts.py index d7dd444..3e72c36 100644 --- a/src/ds_caselaw_utils/test_courts.py +++ b/src/ds_caselaw_utils/test_courts.py @@ -12,6 +12,7 @@ def test_loads_selectable_courts(self): data = [ { "name": "court_group", + "display_name": "court group 1", "courts": [ { "name": "court1", @@ -19,12 +20,23 @@ def test_loads_selectable_courts(self): }, {"name": "court2", "selectable": False}, ], - } + }, + { + "name": "court_group2", + "display_name": "court group 2", + "courts": [{"name": "court3", "selectable": False}], + }, ] repo = CourtsRepository(data) selectable = repo.get_selectable() self.assertIn("court1", [c.name for c in selectable]) self.assertNotIn("court2", [c.name for c in selectable]) + groups = repo.get_selectable_groups() + self.assertIn("court group 1", [g.name for g in groups]) + self.assertNotIn("court group 2", [g.name for g in groups]) + self.assertIn("court1", [c.name for g in groups for c in g.courts]) + self.assertNotIn("court2", [c.name for g in groups for c in g.courts]) + self.assertNotIn("court3", [c.name for g in groups for c in g.courts]) def test_loads_listable_courts(self): data = [