Source code for pudl.extract.eia191

"""Extract EIA Form 191 data from CSVs."""

import pandas as pd
from dagster import Output, asset

from pudl.extract.csv import CsvExtractor
from pudl.extract.extractor import GenericMetadata, PartitionSelection, raw_df_factory




class Extractor(CsvExtractor):
    """Extractor for EIA form 191."""

    def __init__(self, *args, **kwargs):
        """Initialize the module.

        Args:
            ds (:class:datastore.Datastore): Initialized datastore.
        """


        self.METADATA = GenericMetadata("eia191")

        super().__init__(*args, **kwargs)



    def source_filename(self, page: str, **partition: PartitionSelection) -> str:
        """Override filename selection.

        Since archive version 33.0.0, EIA-191 archives include both annual and monthly files for some years.
        All files specify whether they are annual or monthly.
        PUDL processes only monthly EIA-191 data.
        """
        return (
            super().source_filename(page, **partition).replace(".csv", "_monthly.csv")
        )




    def get_page_cols(self, page: str, partition_key: str) -> list[str]:
        """Get the columns for a particular page and partition key.

        EIA 191 data has the same set of columns for all years,
        so regardless of the partition key provided we select the same columns here.
        """
        return super().get_page_cols(page, "any_year")




    def process_raw(
        self, df: pd.DataFrame, page: str, **partition: PartitionSelection
    ) -> pd.DataFrame:
        """Rename columns using ``any_year`` partition."""
        return df.rename(
            columns=self._metadata.get_column_map(page=page, year="any_year")
        )






raw_eia191__all_dfs = raw_df_factory(Extractor, name="eia191")



@asset


def raw_eia191__data(raw_eia191__all_dfs):
    """Extract raw EIA company data from CSV sheets into dataframes.

    Returns:
        An extracted EIA 191 dataframe.
    """
    return Output(value=raw_eia191__all_dfs["data"])