diff --git a/apps/establishment/models.py b/apps/establishment/models.py index 555b26e1..a4e0afa8 100644 --- a/apps/establishment/models.py +++ b/apps/establishment/models.py @@ -2,7 +2,7 @@ from datetime import datetime from functools import reduce from operator import or_ -from typing import List, Union +from typing import List import elasticsearch_dsl from django.conf import settings @@ -184,10 +184,9 @@ class EstablishmentQuerySet(models.QuerySet): """Return establishments by country code""" return self.filter(address__city__country=country) - def by_country_code(self, codes: Union[iter, str]): + def by_country_code(self, code: str): """Return establishments by country code""" - codes = codes if hasattr(codes, '__iter__') else [codes] - return self.filter(address__city__country__code__in=codes) + return self.filter(address__city__country__code=code) def published(self): """