Gleamviz data reader

gleamviz_data.py

import os
from typing import List, Optional

import pandas


class DataReader:
    def __init__(self, root_path: str):
        self.root_path = root_path

    def read_file(self, relative_path: str):
        path = os.path.join(self.root_path, relative_path)
        data = pandas.read_csv(path, sep="\t")
        return data

    @property
    def compartments(self) -> List[str]:
        # TODO: Replace with parsed XML
        return [
            "Recovered",
            "Susceptible",
            "Infectious",
            "Exposed",
        ]

    @property
    def cities(self):
        ...

    @property
    def continents(self):
        ...

    @property
    def countries(self):
        ...

    @property
    def hemispheres(self):
        ...

    @property
    def regions(self):
        ...

    def get_data(self, *, 
        compartment_id: int,
        
        city_id: Optional[int] = None,
        continent_id: Optional[int] = None,
        country_id: Optional[int] = None,
        hemisphere_id: Optional[int] = None,
        region_id: Optional[int] = None
    ) -> pandas.DataFrame:
        if city_id:
            relative_path = f"cities/{city_id}-{compartment_id}.tsv"
        elif continent_id:
            relative_path = f"regions/{continent_id}-{compartment_id}.tsv"
        elif country_id:
            relative_path = f"countries/{country_id}-{compartment_id}.tsv"
        elif hemisphere_id:
            relative_path = f"hemispheres/{hemisphere_id}-{compartment_id}.tsv"
        elif region_id:
            relative_path = f"regions/{region_id}-{compartment_id}.tsv"
        else:
            relative_path = f"global/0-{compartment_id}.tsv"

        # TODO: Check invalid argument combinations

        return self.read_file(relative_path).assign(compartment=self.compartments[compartment_id])