Seperate DataSources to their own class

Author: Jeeven Dhanoa

app/data/__init__.py

@@ -3,19 +3,36 @@
 from ..services.location.jhu import JhuLocationService
 from ..services.location.nyt import NYTLocationService
 
-# Mapping of services to data-sources.
-DATA_SOURCES = {
-    "jhu": JhuLocationService(),
-    "csbs": CSBSLocationService(),
-    "nyt": NYTLocationService(),
-}
 
-
-def data_source(source):
+class DataSources:
     """
-    Retrieves the provided data-source service.
-
-    :returns: The service.
-    :rtype: LocationService
+    Class to represent the root of the aggregate containing the location services.
     """
-    return DATA_SOURCES.get(source.lower())
+
+    # Mapping of services to data-sources.
+    __DATA_SOURCES_MAP = {
+        "jhu": JhuLocationService(),
+        "csbs": CSBSLocationService(),
+        "nyt": NYTLocationService(),
+    }
+
+    def __init__(self):
+        pass
+
+    def get_data_source(self, source):
+        """
+        Retrieves the provided data-source service.
+
+        :returns: The service.
+        :rtype: LocationService
+        """
+        return self.__DATA_SOURCES_MAP.get(source.lower())
+
+    def get_data_sources(self):
+        """
+            Retrieves a dict of all data sources.
+
+            :returns: The dictionary of data sources.
+            :rtype: dict
+        """
+        return self.__DATA_SOURCES_MAP

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 DataSources
 from .routers import V1, V2
 from .utils.httputils import setup_client_session, teardown_client_session
 
@@ -41,6 +41,8 @@
     on_shutdown=[teardown_client_session],
 )
 
+DATA_SOURCES = DataSources()
+
 # #####################
 # Middleware
 #######################
@@ -73,8 +75,8 @@ 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"))
+    # Retrieve the data source from query param.
+    source = DATA_SOURCES.get_data_source(request.query_params.get("source", default="jhu"))
 
     # Abort with 404 if source cannot be found.
     if not source:

app/routers/v2.py

@@ -3,10 +3,11 @@
 
 from fastapi import APIRouter, HTTPException, Request
 
-from ..data import DATA_SOURCES
+from ..data import DataSources
 from ..models import LatestResponse, LocationResponse, LocationsResponse
 
 V2 = APIRouter()
+DATA_SOURCES = DataSources()
 
 
 class Sources(str, enum.Enum):
@@ -107,4 +108,4 @@ async def sources():
     """
     Retrieves a list of data-sources that are availble to use.
     """
-    return {"sources": list(DATA_SOURCES.keys())}
+    return {"sources": list(DATA_SOURCES.get_data_sources().keys())}