From b3ed21252d87384c25f7aa04717ec74c33c517e9 Mon Sep 17 00:00:00 2001 From: Leo Hemsted Date: Thu, 9 Jul 2020 19:10:51 +0100 Subject: [PATCH] ensure you can get broadcast_areas from a list as well as from *args --- notifications_utils/broadcast_areas/__init__.py | 4 ++++ tests/test_broadcast_area.py | 14 ++++++++++++++ 2 files changed, 18 insertions(+) diff --git a/notifications_utils/broadcast_areas/__init__.py b/notifications_utils/broadcast_areas/__init__.py index 493acc4b4..fdc1fd37a 100644 --- a/notifications_utils/broadcast_areas/__init__.py +++ b/notifications_utils/broadcast_areas/__init__.py @@ -147,6 +147,10 @@ class BroadcastAreaLibraries(SerialisedModelCollection, GetItemByIdMixin): yield area def get_areas(self, *area_ids): + # allow people to call `get_areas('a', 'b') or get_areas(['a', 'b'])` + if len(area_ids) == 1 and isinstance(area_ids[0], list): + area_ids = area_ids[0] + return list(itertools.chain(*( [area for area in self.all_areas if area.id == area_id] for area_id in area_ids diff --git a/tests/test_broadcast_area.py b/tests/test_broadcast_area.py index 7f71cb56b..7ef23491d 100644 --- a/tests/test_broadcast_area.py +++ b/tests/test_broadcast_area.py @@ -100,6 +100,16 @@ def test_get_names_of_areas(): ] +def test_get_areas_accepts_lists(): + areas_from_list = broadcast_area_libraries.get_areas( + ['wales', 'vale-of-glamorgan', 'england', 'essex'] + ) + areas_from_args = broadcast_area_libraries.get_areas( + 'wales', 'vale-of-glamorgan', 'england', 'essex', + ) + assert areas_from_args == areas_from_list + + def test_has_polygons(): assert len( @@ -114,6 +124,10 @@ def test_has_polygons(): broadcast_area_libraries.get_polygons_for_areas_long_lat('england', 'scotland') ) == 35 + 195 == 230 + assert len( + broadcast_area_libraries.get_polygons_for_areas_long_lat(['england', 'scotland']) + ) == 35 + 195 == 230 + assert broadcast_area_libraries.get_polygons_for_areas_lat_long('england')[0][0] == [ 55.811085, -2.034358 # https://goo.gl/maps/wsf2LUWzYinwydMk8 ]