Source code for airbase.summary.db

from __future__ import annotations

import sqlite3
import sys
from collections import defaultdict
from contextlib import closing, contextmanager
from pathlib import Path
from typing import Iterator

if sys.version_info >= (3, 11):  # pragma: no cover
    from importlib import resources
else:  # pragma: no cover
    import importlib_resources as resources


def summary() -> Path:
    source = resources.files(__package__) / "summary.sqlite"
    path: Path
    with resources.as_file(source) as path:
        return path


[docs]class DB: """ In DB containing the available country and pollutants """ db = sqlite3.connect(f"file:{summary()}?mode=ro", uri=True)
[docs] @classmethod @contextmanager def cursor(cls) -> Iterator[sqlite3.Cursor]: """db cursor as a "self closing" context manager""" with closing(cls.db.cursor()) as cur: yield cur
[docs] @classmethod def countries(cls) -> list[str]: """ Get the list of unique countries from the summary. :return: list of available country codes """ with cls.cursor() as cur: cur.execute("SELECT country_code FROM countries;") return list(row[0] for row in cur.fetchall())
[docs] @classmethod def pollutants(cls) -> dict[str, str]: """ Get the list of unique pollutants from the summary. :param summary: The E1a summary. :return: The available pollutants, as a dictionary with with name as keys with name as values, e.g. {"NO": "38", ...} """ with cls.cursor() as cur: cur.execute("SELECT pollutant, pollutant_id FROM pollutants;") return dict(cur.fetchall())
[docs] @classmethod def search_pollutant( cls, query: str, *, limit: int | None = None ) -> dict[str, int]: """ Search for a pollutant's ID number based on its name. :param query: The pollutant to search for. :param limit: (optional) Max number of results. :return: The best pollutant matches, as a dictionary with with name as keys with name as values, e.g. {"NO": 38, ...} """ with cls.cursor() as cur: cur.execute( f""" SELECT pollutant, pollutant_id FROM pollutants WHERE pollutant LIKE '%{query}%' {f"LIMIT {limit}" if limit else ""}; """ ) return dict(cur.fetchall())
[docs] @classmethod def pollutants_per_country(cls) -> dict[str, dict[str, int]]: """ Get the available pollutants per country from the summary. :return: All available pollutants per country, as a dictionary with with country code as keys and a dictionary of pollutant/ids (e.g. {"NO": 38, ...}) as values. """ with cls.cursor() as cur: cur.execute( "SELECT country_code, pollutant, pollutant_id FROM summary" ) output: dict[str, dict[str, int]] = defaultdict(dict) for country_code, pollutant, pollutant_id in cur: output[country_code][pollutant] = pollutant_id return dict(output)