Source code for pudl.extract.sec10k

"""Load pre-processed SEC 10-K assets from Google Cloud Storage.

These "raw" tables are generated by the SEC 10-K data extraction pipeline which can be
found in this repository: https://github.com/catalyst-cooperative/mozilla-sec-eia

Upstream data is not partitioned by year, but we want to be able to extract a subset of
the data for testing, so the :class:`Sec10kDataConfig` allow specification of which years
to extract, and those are used to filter the extracted data before returning it.
"""

from io import BytesIO

import dagster as dg
import pandas as pd

from pudl.settings import DataSource, Sec10kDataConfig
from pudl.workspace.datastore import Datastore




def extract(ds: Datastore, table: str, years: list[int]) -> pd.DataFrame:
    """Extract SEC 10-K data from the datastore.

    Allows filtering by year to enable testing of the pipeline with a smaller amount
    of data, like a pseudo-partition. This is necessary because the SEC 10-K data is
    not partitioned upstream.

    Args:
        ds: Initialized PUDL datastore.
        table: Which of the valid tables should be extracted?
        years: Which years of data to include in the output.

    Returns:
        A dataframe containing the SEC 10-K data.
    """
    resource = ds.get_unique_resource("sec10k", table=table)
    df = pd.read_parquet(BytesIO(resource))
    if "year_quarter" in df.columns:
        date_col = "year_quarter"
    elif "report_date" in df.columns:
        date_col = "report_date"
    else:  # pragma: no cover
        raise ValueError(
            f"Expected either year_quarter or report_date in {table}."
            "but found neither. Cannot identify reporting year."
        )
    # Filter the dataframe for years specified by the data config:
    years_mask = df[date_col].astype("string").str[:4].astype(int).isin(years)
    return df.loc[years_mask]





def raw_sec10k_asset_factory(table) -> dg.AssetsDefinition:
    """An asset factory for extracting SEC 10-K data by table."""

    @dg.asset(
        group_name="raw_sec10k",
        name=table,
        required_resource_keys={"datastore", "global_data_config"},
    )
    def sec10k_asset(context) -> pd.DataFrame:
        sec10k_data_config: Sec10kDataConfig = (
            context.resources.global_data_config.pudl.sec10k
        )
        ds: Datastore = context.resources.datastore
        if table in sec10k_data_config.tables:
            return extract(ds=ds, table=table, years=sec10k_data_config.years)
        raise ValueError(f"Unrecognized SEC 10-K table: {table}")  # pragma: no cover

    return sec10k_asset





raw_sec10k_assets = [
    raw_sec10k_asset_factory(table)
    for table in DataSource.from_id("sec10k").working_partitions["tables"]
]