src/customer_intelligence/etl/extract.py

"""Extract raw Olist CSV files into DataFrames."""

from pathlib import Path

import pandas as pd


class OlistExtractor:
    """Reads each Olist CSV file from the raw data directory."""

    CSV_MAP = {
        "orders": "olist_orders_dataset.csv",
        "order_items": "olist_order_items_dataset.csv",
        "order_payments": "olist_order_payments_dataset.csv",
        "order_reviews": "olist_order_reviews_dataset.csv",
        "customers": "olist_customers_dataset.csv",
        "sellers": "olist_sellers_dataset.csv",
        "products": "olist_products_dataset.csv",
        "category_translations": "product_category_name_translation.csv",
        "geolocation": "olist_geolocation_dataset.csv",
    }

    def __init__(self, raw_dir: str | Path) -> None:
        self.raw_dir = Path(raw_dir)

    def _read(self, key: str, **kwargs) -> pd.DataFrame:
        path = self.raw_dir / self.CSV_MAP[key]
        if not path.exists():
            raise FileNotFoundError(f"Expected CSV not found: {path}")
        return pd.read_csv(path, dtype=str, keep_default_na=False, **kwargs)

    def extract_orders(self) -> pd.DataFrame:
        return self._read("orders")

    def extract_order_items(self) -> pd.DataFrame:
        return self._read("order_items")

    def extract_order_payments(self) -> pd.DataFrame:
        return self._read("order_payments")

    def extract_order_reviews(self) -> pd.DataFrame:
        return self._read("order_reviews")

    def extract_customers(self) -> pd.DataFrame:
        return self._read("customers")

    def extract_sellers(self) -> pd.DataFrame:
        return self._read("sellers")

    def extract_products(self) -> pd.DataFrame:
        return self._read("products")

    def extract_category_translations(self) -> pd.DataFrame:
        return self._read("category_translations")

    def extract_geolocation(self) -> pd.DataFrame:
        return self._read("geolocation")

    def extract_all(self) -> dict[str, pd.DataFrame]:
        return {
            "orders": self.extract_orders(),
            "order_items": self.extract_order_items(),
            "order_payments": self.extract_order_payments(),
            "order_reviews": self.extract_order_reviews(),
            "customers": self.extract_customers(),
            "sellers": self.extract_sellers(),
            "products": self.extract_products(),
            "category_translations": self.extract_category_translations(),
            "geolocation": self.extract_geolocation(),
        }