Refactoring

app/enums/sources.py

@@ -0,0 +1,8 @@
+from enum import Enum
+
+class Sources(str, Enum):
+    """
+    A source available for retrieving data.
+    """
+    jhu = 'jhu'
+    csbs = 'csbs'

app/main.py

@@ -1,42 +1,34 @@
 """
 app.main.py
 """
-import datetime as dt
-import enum
 import logging
 import os
 import reprlib
-from typing import Dict, List
+import datetime as dt
 
-import fastapi
 import pydantic
 import uvicorn
-from fastapi.middleware.wsgi import WSGIMiddleware
-from fastapi.middleware.cors import CORSMiddleware
 
-from . import models
-from .core import create_app
-from .data import data_source, data_sources
+from fastapi import FastAPI
+from fastapi import Request, Response
 
-# ################
-# Dependencies
-# ################
+from fastapi.responses import JSONResponse
 
+from fastapi.middleware.wsgi import WSGIMiddleware
+from fastapi.middleware.cors import CORSMiddleware
 
-class Sources(str, enum.Enum):
-    """
-    A source available for retrieving data.
-    """
-    jhu = 'jhu'
-    csbs = 'csbs'
+from .core import create_app
+from .data import data_source
 
+from .models.location import LocationResponse as Location, LocationsResponse as Locations
+from .models.latest import LatestResponse as Latest
 
 # ############
 # FastAPI App
 # ############
 LOGGER = logging.getLogger('api')
 
-APP = fastapi.FastAPI(
+APP = FastAPI(
     title='Coronavirus Tracker',
     description='API for tracking the global coronavirus (COVID-19, SARS-CoV-2) outbreak. Project page: https://github.com/ExpDev07/coronavirus-tracker-api.',
     version='2.0.1',
@@ -59,7 +51,7 @@ class Sources(str, enum.Enum):
 
 # TODO this could probably just be a FastAPI dependency.
 @APP.middleware('http')
-async def add_datasource(request: fastapi.Request, call_next):
+async def add_datasource(request: Request, call_next):
     """
     Attach the data source to the request.state.
     """
@@ -68,7 +60,7 @@ async def add_datasource(request: fastapi.Request, call_next):
 
     # Abort with 404 if source cannot be found.
     if not source:
-        return fastapi.Response('The provided data-source was not found.', status_code=404)
+        return Response('The provided data-source was not found.', status_code=404)
 
     # Attach source to request.
     request.state.source = source
@@ -86,104 +78,22 @@ async def add_datasource(request: fastapi.Request, call_next):
 
 @APP.exception_handler(pydantic.error_wrappers.ValidationError)
 async def handle_validation_error(
-    request: fastapi.Request, exc: pydantic.error_wrappers.ValidationError
+    request: Request, exc: pydantic.error_wrappers.ValidationError
 ):
     """
     Handles validation errors.
     """
-    return fastapi.responses.JSONResponse({'message': exc.errors()}, status_code=422)
+    return JSONResponse({'message': exc.errors()}, status_code=422)
 
 
 # ################
-# Routes
+# Routing
 # ################
 
-V2 = fastapi.APIRouter()
-
-
-@V2.get('/latest', response_model=models.LatestResponse)
-def get_latest(request: fastapi.Request, source: Sources = 'jhu'):
-    """
-    Getting latest amount of total confirmed cases, deaths, and recoveries.
-    """
-    locations = request.state.source.get_all()
-    return {
-        'latest': {
-            'confirmed': sum(map(lambda location: location.confirmed, locations)),
-            'deaths'   : sum(map(lambda location: location.deaths, locations)),
-            'recovered': sum(map(lambda location: location.recovered, locations)),
-        }
-    }
-
-
-@V2.get(
-    '/locations', response_model=models.LocationsResponse, response_model_exclude_unset=True
-)
-def get_locations(
-    request: fastapi.Request,
-    source: Sources = 'jhu',
-    country_code: str = None,
-    province: str = None,
-    county: str = None,
-    timelines: bool = False,
-):
-    """
-    Getting the locations.
-    """
-    # All query paramameters.
-    params = dict(request.query_params)
-
-    # Remove reserved params.
-    params.pop('source', None)
-    params.pop('timelines', None)
-
-    # Retrieve all the locations.
-    locations = request.state.source.get_all()
-
-    # Attempt to filter out locations with properties matching the provided query params.
-    for key, value in params.items():
-        # Clean keys for security purposes.
-        key   = key.lower()
-        value = value.lower().strip('__')
-
-        # Do filtering.
-        try:
-            locations = [location for location in locations if str(getattr(location, key)).lower() == str(value)]
-        except AttributeError:
-            pass
-
-    # Return final serialized data.
-    return {
-        'latest': {
-            'confirmed': sum(map(lambda location: location.confirmed, locations)),
-            'deaths'   : sum(map(lambda location: location.deaths, locations)),
-            'recovered': sum(map(lambda location: location.recovered, locations)),
-        },
-        'locations': [location.serialize(timelines) for location in locations],
-    }
-
-
-@V2.get('/locations/{id}', response_model=models.LocationResponse)
-def get_location_by_id(request: fastapi.Request, id: int, source: Sources = 'jhu', timelines: bool = True):
-    """
-    Getting specific location by id.
-    """
-    return {
-        'location': request.state.source.get(id).serialize(timelines)
-    }
-
-
-@V2.get('/sources')
-async def sources():
-    """
-    Retrieves a list of data-sources that are availble to use.
-    """
-    return {
-        'sources': list(data_sources.keys())
-    }
+from .router import router
 
 # Include routers.
-APP.include_router(V2, prefix='/v2', tags=['v2'])
+APP.include_router(router, prefix='/v2', tags=['v2'])
 
 # mount the existing Flask app
 # v1 @ /

app/models/latest.py

@@ -0,0 +1,15 @@
+from pydantic import BaseModel
+
+class Latest(BaseModel):
+    """
+    Latest model.
+    """
+    confirmed: int
+    deaths: int
+    recovered: int
+
+class LatestResponse(BaseModel):
+    """
+    Response for latest.
+    """
+    latest: Latest

app/models/location.py

@@ -0,0 +1,31 @@
+from pydantic import BaseModel
+from typing import Dict, List
+from .timeline import Timelines
+from .latest import Latest
+
+class Location(BaseModel):
+    """
+    Location model.
+    """
+    id: int
+    country: str
+    country_code: str
+    county: str = ''
+    province: str = ''
+    last_updated: str # TODO use datetime.datetime type.
+    coordinates: Dict
+    latest: Latest
+    timelines: Timelines = {}
+
+class LocationResponse(BaseModel):
+    """
+    Response for location.
+    """
+    location: Location
+
+class LocationsResponse(BaseModel):
+    """
+    Response for locations.
+    """
+    latest: Latest
+    locations: List[Location] = []

app/models/timeline.py

@@ -0,0 +1,18 @@
+from pydantic import BaseModel
+from typing import Dict
+
+class Timeline(BaseModel):
+    """
+    Timeline model.
+    """
+
+    latest: int
+    timeline: Dict[str, int] = {}
+
+class Timelines(BaseModel):
+    """
+    Timelines model.
+    """
+    confirmed: Timeline
+    deaths: Timeline
+    recovered: Timeline

app/router/__init__.py

@@ -0,0 +1,7 @@
+from fastapi import APIRouter
+
+# Create the router.
+router = APIRouter()
+
+# The routes.
+from . import latest, sources, locations

app/router/latest.py

@@ -0,0 +1,18 @@
+from fastapi import Request
+from . import router
+from ..enums.sources import Sources
+from ..models.latest import LatestResponse as Latest
+
+@router.get('/latest', response_model=Latest)
+def get_latest(request: Request, source: Sources = 'jhu'):
+    """
+    Getting latest amount of total confirmed cases, deaths, and recoveries.
+    """
+    locations = request.state.source.get_all()
+    return {
+        'latest': {
+            'confirmed': sum(map(lambda location: location.confirmed, locations)),
+            'deaths'   : sum(map(lambda location: location.deaths, locations)),
+            'recovered': sum(map(lambda location: location.recovered, locations)),
+        }
+    }