ETL

etl.py

from typing import Dict, Any, Optional


def extract(raw_data: Dict[Any, Any]) -> Dict[Any, Any]:
    if not isinstance(raw_data, dict):
        raise TypeError("Failed to load data: Data type must be dict")

    return raw_data


def _change_data_types(raw_data: Dict[Any, Any]) -> Dict[str, int]:
    try:
        return {str(k): int(v) for k, v in raw_data.items()}
    except Exception as e:
        raise e


def transform(struct: Dict[str, int]) -> Dict[str, int]:
    prepared_data: Dict[str, int] = _change_data_types(struct)
    return {k: v ** 2 for k, v in prepared_data.items()}


def load(
    struct: Dict[str, int], database_id: str
) -> Optional[Dict[str, Dict[str, int]]]:
    database: Dict[str, Dict[str, int]] = {"squared": {}}
    if database_id == "squared":
        database[database_id] = struct
        return database

    return None


if __name__ == "__main__":
    # Define the data
    data: Dict[str, int] = {
        "a": 213,
        "b": 314,
    }
    _id: str = "squared"

    # ETL
    new_data: Optional[Dict[str, Dict[str, int]]] = load(transform(extract(data)), _id)
    if new_data:
        print(new_data)  # {'squared': {'a': 45369, 'b': 98596}}
        assert new_data == {"squared": {"a": 45369, "b": 98596}}