ahmet / Python-for-Data-Science

import matplotlib.pyplot as plt
import matplotlib.ticker as ticker
from load_csv import load


def to_number(value) -> float:
    """Convert a Gapminder value such as '4.83M' into a float."""
    suffixes = {"k": 1e3, "m": 1e6, "b": 1e9}
    if isinstance(value, str):
        value = value.strip()
        if value and value[-1].lower() in suffixes:
            return float(value[:-1]) * suffixes[value[-1].lower()]
    return float(value)


def format_million(value, _position) -> str:
    """Format an axis tick value as a number of millions."""
    return f"{value / 1e6:g}M"


def population_of(dataset, country: str, years: list) -> list:
    """Return the population of a country for the given years.

    None is returned when the country is missing or holds bad values.
    """
    row = dataset[dataset["country"] == country]
    if row.empty:
        print(f"Error: country not found: {country}")
        return None
    try:
        return [to_number(value) for value in row.iloc[0][years]]
    except ValueError:
        print(f"Error: unexpected values for {country}")
        return None


def main() -> None:
    """Compare the population of my campus country with another one."""
    countries = ["Turkey", "France"]
    dataset = load("../population_total.csv")
    if dataset is None:
        return
    if "country" not in dataset.columns:
        print("Error: the dataset has no 'country' column")
        return
    try:
        years = [year for year in dataset.columns[1:]
                 if 1800 <= int(year) <= 2050]
    except ValueError:
        print("Error: the dataset has unexpected column names")
        return
    for country in countries:
        values = population_of(dataset, country, years)
        if values is None:
            return
        plt.plot([int(year) for year in years], values, label=country)
    plt.title("Population Projections")
    plt.xlabel("Year")
    plt.ylabel("Population")
    plt.gca().yaxis.set_major_formatter(
        ticker.FuncFormatter(format_million))
    plt.legend()
    plt.show()


if __name__ == "__main__":
    main()