use Source enum instead of magic strings

Author: Kilo59

app/data/__init__.py

@@ -18,9 +18,9 @@ class Sources(str, enum.Enum):
 
 # Mapping of services to data-sources.
 DATA_SOURCES = {
-    "jhu": JhuLocationService(),
-    "csbs": CSBSLocationService(),
-    "nyt": NYTLocationService(),
+    Sources.jhu: JhuLocationService(),
+    Sources.csbs: CSBSLocationService(),
+    Sources.nyt: NYTLocationService(),
 }
 
 

app/main.py

@@ -14,7 +14,7 @@
 from sentry_sdk.integrations.asgi import SentryAsgiMiddleware
 
 from .config import get_settings
-from .data import data_source
+from .data import Sources, data_source
 from .routers import V1, V2
 from .utils.httputils import setup_client_session, teardown_client_session
 
@@ -74,7 +74,7 @@ async def add_datasource(request: Request, call_next):
     Attach the data source to the request.state.
     """
     # Retrieve the datas ource from query param.
-    source = data_source(request.query_params.get("source", default="jhu"))
+    source = data_source(request.query_params.get("source", default=Sources.jhu))
 
     # Abort with 404 if source cannot be found.
     if not source:

app/routers/v2.py

@@ -10,7 +10,9 @@
 
 
 @V2.get("/latest", response_model=LatestResponse)
-async def get_latest(request: Request, source: Sources = "jhu"):  # pylint: disable=unused-argument
+async def get_latest(
+    request: Request, source: Sources = Sources.jhu
+):  # pylint: disable=unused-argument
     """
     Getting latest amount of total confirmed cases, deaths, and recoveries.
     """
@@ -28,7 +30,7 @@ async def get_latest(request: Request, source: Sources = "jhu"):  # pylint: disa
 @V2.get("/locations", response_model=LocationsResponse, response_model_exclude_unset=True)
 async def get_locations(
     request: Request,
-    source: Sources = "jhu",
+    source: Sources = Sources.jhu,
     country_code: str = None,
     province: str = None,
     county: str = None,
@@ -81,7 +83,7 @@ async def get_locations(
 # pylint: disable=invalid-name
 @V2.get("/locations/{id}", response_model=LocationResponse)
 async def get_location_by_id(
-    request: Request, id: int, source: Sources = "jhu", timelines: bool = True
+    request: Request, id: int, source: Sources = Sources.jhu, timelines: bool = True
 ):
     """
     Getting specific location by id.