ahmet / Python-for-Data-Science

import matplotlib.pyplot as plt
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 merge_year(income, life, year: str):
    """Join both datasets on the country for a single year.

    None is returned when the year is missing or holds bad values.
    """
    for dataset in (income, life):
        if "country" not in dataset.columns or year not in dataset.columns:
            print(f"Error: the datasets must hold a country and a {year}")
            return None
    merged = income[["country", year]].merge(
        life[["country", year]], on="country", suffixes=("_gdp", "_life")
    )
    merged = merged.dropna()
    if merged.empty:
        print(f"Error: no country holds data for {year}")
        return None
    try:
        gdp = [to_number(value) for value in merged[f"{year}_gdp"]]
        expectancy = [float(value) for value in merged[f"{year}_life"]]
    except ValueError:
        print(f"Error: unexpected values for the year {year}")
        return None
    return gdp, expectancy


def main() -> None:
    """Plot life expectancy against income per person for 1900."""
    year = "1900"
    income_file = ("../income_per_person_gdppercapita_ppp"
                   "_inflation_adjusted.csv")
    income = load(income_file)
    life = load("../life_expectancy_years.csv")
    if income is None or life is None:
        return
    data = merge_year(income, life, year)
    if data is None:
        return
    gdp, expectancy = data
    plt.scatter(gdp, expectancy, label=f"Countries in {year}")
    plt.xscale("log")
    plt.xticks([300, 1000, 10000], ["300", "1k", "10k"])
    plt.title(year)
    plt.xlabel("Gross domestic product")
    plt.ylabel("Life Expectancy")
    plt.legend()
    plt.show()


if __name__ == "__main__":
    main()