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)