00_core.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# hide\n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp core"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# lockdowndates\n",
    "\n",
    "> Retrieve the dates of restructions imposed on countries around the world during the covid pandemic. Helpful for maching learning projects with a time element during the feature engineering phase."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#hide\n",
    "from nbdev.showdoc import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "from datetime import datetime as dt\n",
    "from typing import List, Union, Tuple, Optional"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "class LockdownDates:\n",
    "    '''\n",
    "       Retrieve the dates of the restrictions in countries imposed by governments around the world during the covid-19 pandemic. \n",
    "       \n",
    "       `country`: Country from table of countries in README.md\n",
    "       <br/>`start_date`: Date you wish to collect dates from in \"YYYY-MM-DD\" format\n",
    "       <br/>`end_date`: Date you wish to collect dates from in \"YYYY-MM-DD\" format\n",
    "       `restrictions`: List of restrictions to be returned listed in README.md\n",
    "    '''\n",
    "        \n",
    "    def __init__(self, country:Union[List[str],str], start_date:str, end_date:str, restrictions: Union[List, Tuple]):\n",
    "        \n",
    "        if isinstance(country, list):\n",
    "            self.country = country\n",
    "        else:\n",
    "            self.country = [country]\n",
    "            \n",
    "        if isinstance(start_date, str):\n",
    "            self.start_date = dt.strptime(start_date, \"%Y-%m-%d\")\n",
    "        else:\n",
    "            print(\"Incorrect format for start_date, expecting %Y-%m-%d\")\n",
    "            \n",
    "        if isinstance(end_date, str):\n",
    "            self.end_date = dt.strptime(end_date, \"%Y-%m-%d\")\n",
    "        else:\n",
    "            print(\"Incorrect format for end_date, expecting %Y-%m-%d\")\n",
    "            \n",
    "        if isinstance(restrictions, list) or isinstance(restrictions, tuple):\n",
    "            self.restrictions = restrictions\n",
    "            self.restriction_keys = {\n",
    "                \"stay_at_home\" : \"C6_Stay at home requirements\",\n",
    "                \"masks\" : \"H6_Facial Coverings\"\n",
    "            }\n",
    "        else:\n",
    "            print(f\"Incorrect format for restrictions, you provided {type(restrictions)}, needed to be list or tuple\")\n",
    "        \n",
    "    def fetch(self) -> pd.DataFrame:\n",
    "\n",
    "        restrictions_to_be_used = [self.restriction_keys[restriction] for restriction in self.restrictions]\n",
    "        \n",
    "        usecols= [\"CountryName\", \"CountryCode\", \"Date\"] + restrictions_to_be_used\n",
    "        \n",
    "        df_dtype = {\n",
    "            \"CountryName\": str,\n",
    "            \"CountryCode\": str,\n",
    "            \"Date\": str,\n",
    "        }\n",
    "\n",
    "        print(\"Fetching lockdown dates...\")\n",
    "        try:\n",
    "            urls = [f\"https://github.com/seanyboi/lockdowndates_data/blob/main/data/{country.lower().replace(' ', '')}.parquet?raw=true\" for country in self.country]\n",
    "            lockdown_df = pd.concat((pd.read_parquet(u, columns=usecols, engine=\"pyarrow\") for u in urls))\n",
    "            lockdown_df = lockdown_df.astype(df_dtype)\n",
    "            return lockdown_df\n",
    "        except Exception as e:\n",
    "            print(f\"Error fetching lockdown data - {e}\")\n",
    "\n",
    "        \n",
    "    def engineer_df(self) -> pd.DataFrame:\n",
    "        # fetch data\n",
    "        df = self.fetch()\n",
    "        \n",
    "        try:\n",
    "            \n",
    "            restriction_names = dict(zip(self.restriction_keys.values(), self.restriction_keys.keys()))\n",
    "            # rename columns\n",
    "            rename_columns = columns={\n",
    "                \"CountryName\": \"country\", \n",
    "                \"CountryCode\": \"country_code\", \n",
    "                \"Date\": \"timestamp\", \n",
    "                **restriction_names\n",
    "            }\n",
    "\n",
    "            df = df.rename(columns=rename_columns)\n",
    "\n",
    "            # configure dates and set_index\n",
    "            df[\"timestamp\"] = df[\"timestamp\"].str.replace(r'(\\d{4})(\\d{2})(\\d{2})', r'\\g<1>-\\g<2>-\\g<3>', regex=True)\n",
    "            df[\"timestamp\"] = pd.to_datetime(df[\"timestamp\"])\n",
    "\n",
    "            # convert columns to categories\n",
    "            restriction_cats = list(self.restriction_keys.keys()) + ['country', 'country_code']\n",
    "            for col in restriction_cats:\n",
    "                df[col] = df[col].astype('category')\n",
    "\n",
    "            if len(self.country) == 1:\n",
    "                print(f\"Fetched lockdown dates for: {self.country[0]}\")\n",
    "            else:\n",
    "                print(f\"Fetched lockdown dates for: {', '.join(self.country)}\")\n",
    "\n",
    "            return df\n",
    "        except Exception as e:\n",
    "            print(f\"Formatting data failed - please raise an issue on our repo! - {e}\")\n",
    "    \n",
    "    def filter_df(self) -> pd.DataFrame:\n",
    "        df = self.engineer_df()\n",
    "        try: \n",
    "            df = df[df['country'].isin(self.country)]\n",
    "            df = df.pivot_table(index=\"timestamp\", columns='country', aggfunc='first')\n",
    "            df.columns = [\"{}_{}\".format(col[1].lower().replace(\"'\", \"\").replace(\" \", \"\"), col[0]) for col in df.columns.values]\n",
    "            df = df.loc[self.start_date : self.end_date]\n",
    "        \n",
    "            if df.empty and len(self.country) == 1:\n",
    "                raise Exception(f\"No lockdown data for {self.country[0]}\")\n",
    "            if df.empty:\n",
    "                raise Exception(f\"No lockdown data for {self.country}\")\n",
    "        except:\n",
    "            print(f\"No lockdown data for {self.country} between {self.start_date} and {self.end_date}\")\n",
    "        \n",
    "        return df\n",
    "    \n",
    "    def dates(self, save:bool = False) -> pd.DataFrame:\n",
    "        '''\n",
    "        Returns the restriction lockdown dates for a specific set of countries.\n",
    "        \n",
    "        <b>Parameters</b>\n",
    "            <br/> &nbsp;&nbsp;&nbsp;&nbsp; `save` : bool, optional\n",
    "            <br/>&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp; saves restrictions to a csv file for caching (default is False)\n",
    "    \n",
    "        <b>Returns</b>:\n",
    "            <br/> &nbsp;&nbsp;&nbsp;&nbsp; DataFrame containing the dates a country was subject to certain restrictions during the pandemic.\n",
    "\n",
    "        <b>Raises</b>:\n",
    "            <br/> &nbsp;&nbsp;&nbsp;&nbsp; Exception: failed to collect data.\n",
    "        '''\n",
    "        restrictions = self.filter_df()\n",
    "        try:\n",
    "            if save:\n",
    "                output_file = f\"{self.country.lower().replace(' ','')}-lockdown-restrictions.csv\"\n",
    "                output_dir = Path(\"lockdown_data\")\n",
    "                output_dir.mkdir(parents=True, exist_ok=True)\n",
    "                restrictions.to_csv(output_dir / output_file, index=False)\n",
    "                print(f\"Saved restrictions to - {output_dir}/{output_file}\")\n",
    "        except Exception as e:\n",
    "            print(f\"Failed to save restrictions to csv file - {e}\")\n",
    "        \n",
    "        return restrictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Fetching lockdown dates...\n",
      "Fetched lockdown dates for: Aruba\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>aruba_country_code</th>\n",
       "      <th>aruba_masks</th>\n",
       "      <th>aruba_stay_at_home</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>timestamp</th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>2022-01-01</th>\n",
       "      <td>ABW</td>\n",
       "      <td>2.0</td>\n",
       "      <td>2.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2022-01-02</th>\n",
       "      <td>ABW</td>\n",
       "      <td>2.0</td>\n",
       "      <td>2.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2022-01-03</th>\n",
       "      <td>ABW</td>\n",
       "      <td>2.0</td>\n",
       "      <td>2.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2022-01-04</th>\n",
       "      <td>ABW</td>\n",
       "      <td>2.0</td>\n",
       "      <td>2.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2022-01-05</th>\n",
       "      <td>ABW</td>\n",
       "      <td>2.0</td>\n",
       "      <td>2.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2022-01-06</th>\n",
       "      <td>ABW</td>\n",
       "      <td>2.0</td>\n",
       "      <td>2.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2022-01-07</th>\n",
       "      <td>ABW</td>\n",
       "      <td>2.0</td>\n",
       "      <td>2.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2022-01-08</th>\n",
       "      <td>ABW</td>\n",
       "      <td>2.0</td>\n",
       "      <td>2.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2022-01-09</th>\n",
       "      <td>ABW</td>\n",
       "      <td>2.0</td>\n",
       "      <td>2.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2022-01-10</th>\n",
       "      <td>ABW</td>\n",
       "      <td>2.0</td>\n",
       "      <td>2.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2022-01-11</th>\n",
       "      <td>ABW</td>\n",
       "      <td>2.0</td>\n",
       "      <td>2.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2022-01-12</th>\n",
       "      <td>ABW</td>\n",
       "      <td>2.0</td>\n",
       "      <td>2.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2022-01-13</th>\n",
       "      <td>ABW</td>\n",
       "      <td>2.0</td>\n",
       "      <td>2.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2022-01-14</th>\n",
       "      <td>ABW</td>\n",
       "      <td>2.0</td>\n",
       "      <td>2.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2022-01-15</th>\n",
       "      <td>ABW</td>\n",
       "      <td>2.0</td>\n",
       "      <td>2.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2022-01-16</th>\n",
       "      <td>ABW</td>\n",
       "      <td>2.0</td>\n",
       "      <td>2.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2022-01-17</th>\n",
       "      <td>ABW</td>\n",
       "      <td>2.0</td>\n",
       "      <td>2.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2022-01-18</th>\n",
       "      <td>ABW</td>\n",
       "      <td>2.0</td>\n",
       "      <td>2.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2022-01-19</th>\n",
       "      <td>ABW</td>\n",
       "      <td>2.0</td>\n",
       "      <td>2.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2022-01-20</th>\n",
       "      <td>ABW</td>\n",
       "      <td>2.0</td>\n",
       "      <td>2.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2022-01-21</th>\n",
       "      <td>ABW</td>\n",
       "      <td>2.0</td>\n",
       "      <td>2.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2022-01-22</th>\n",
       "      <td>ABW</td>\n",
       "      <td>2.0</td>\n",
       "      <td>2.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2022-01-23</th>\n",
       "      <td>ABW</td>\n",
       "      <td>2.0</td>\n",
       "      <td>2.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2022-01-24</th>\n",
       "      <td>ABW</td>\n",
       "      <td>2.0</td>\n",
       "      <td>2.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2022-01-25</th>\n",
       "      <td>ABW</td>\n",
       "      <td>2.0</td>\n",
       "      <td>2.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2022-01-26</th>\n",
       "      <td>ABW</td>\n",
       "      <td>2.0</td>\n",
       "      <td>2.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2022-01-27</th>\n",
       "      <td>ABW</td>\n",
       "      <td>2.0</td>\n",
       "      <td>2.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2022-01-28</th>\n",
       "      <td>ABW</td>\n",
       "      <td>2.0</td>\n",
       "      <td>2.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2022-01-29</th>\n",
       "      <td>ABW</td>\n",
       "      <td>2.0</td>\n",
       "      <td>2.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2022-01-30</th>\n",
       "      <td>ABW</td>\n",
       "      <td>2.0</td>\n",
       "      <td>2.0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "           aruba_country_code aruba_masks aruba_stay_at_home\n",
       "timestamp                                                   \n",
       "2022-01-01                ABW         2.0                2.0\n",
       "2022-01-02                ABW         2.0                2.0\n",
       "2022-01-03                ABW         2.0                2.0\n",
       "2022-01-04                ABW         2.0                2.0\n",
       "2022-01-05                ABW         2.0                2.0\n",
       "2022-01-06                ABW         2.0                2.0\n",
       "2022-01-07                ABW         2.0                2.0\n",
       "2022-01-08                ABW         2.0                2.0\n",
       "2022-01-09                ABW         2.0                2.0\n",
       "2022-01-10                ABW         2.0                2.0\n",
       "2022-01-11                ABW         2.0                2.0\n",
       "2022-01-12                ABW         2.0                2.0\n",
       "2022-01-13                ABW         2.0                2.0\n",
       "2022-01-14                ABW         2.0                2.0\n",
       "2022-01-15                ABW         2.0                2.0\n",
       "2022-01-16                ABW         2.0                2.0\n",
       "2022-01-17                ABW         2.0                2.0\n",
       "2022-01-18                ABW         2.0                2.0\n",
       "2022-01-19                ABW         2.0                2.0\n",
       "2022-01-20                ABW         2.0                2.0\n",
       "2022-01-21                ABW         2.0                2.0\n",
       "2022-01-22                ABW         2.0                2.0\n",
       "2022-01-23                ABW         2.0                2.0\n",
       "2022-01-24                ABW         2.0                2.0\n",
       "2022-01-25                ABW         2.0                2.0\n",
       "2022-01-26                ABW         2.0                2.0\n",
       "2022-01-27                ABW         2.0                2.0\n",
       "2022-01-28                ABW         2.0                2.0\n",
       "2022-01-29                ABW         2.0                2.0\n",
       "2022-01-30                ABW         2.0                2.0"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ld = LockdownDates(\"Aruba\", \"2022-01-01\", \"2022-01-30\", (\"stay_at_home\", \"masks\"))\n",
    "x = ld.dates()\n",
    "x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "<h4 id=\"LockdownDates.dates\" class=\"doc_header\"><code>LockdownDates.dates</code><a href=\"__main__.py#L97\" class=\"source_link\" style=\"float:right\">[source]</a></h4>\n",
       "\n",
       "> <code>LockdownDates.dates</code>(**`save`**:`bool`=*`False`*)\n",
       "\n",
       "Returns the restriction lockdown dates for a specific set of countries.\n",
       "\n",
       "<b>Parameters</b>\n",
       "    <br/> &nbsp;&nbsp;&nbsp;&nbsp; `save` : bool, optional\n",
       "    <br/>&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp; saves restrictions to a csv file for caching (default is False)\n",
       "\n",
       "<b>Returns</b>:\n",
       "    <br/> &nbsp;&nbsp;&nbsp;&nbsp; DataFrame containing the dates a country was subject to certain restrictions during the pandemic.\n",
       "\n",
       "<b>Raises</b>:\n",
       "    <br/> &nbsp;&nbsp;&nbsp;&nbsp; Exception: failed to collect data."
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "show_doc(LockdownDates.dates)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#ignore\n",
    "ld = LockdownDates(\"Aruba\", \"2022-01-01\", \"2022-01-30\", [\"stay_at_home\", \"masks\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#ignore\n",
    "assert ld.country==[\"Aruba\"]\n",
    "assert type(ld.country)==list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#ignore\n",
    "assert ld.start_date==dt(2022, 1, 1, 0, 0)\n",
    "assert type(ld.start_date)==dt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#ignore\n",
    "assert ld.end_date==dt(2022, 1, 30, 0, 0)\n",
    "assert type(ld.end_date)==dt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#ignore\n",
    "assert ld.restrictions==[\"stay_at_home\", \"masks\"]\n",
    "assert type(ld.restrictions)==list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Converted 00_core.ipynb.\n",
      "Converted index.ipynb.\n"
     ]
    }
   ],
   "source": [
    "# hide\n",
    "from nbdev.export import notebook2script; notebook2script()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}