Python 用于构建“pandas.DataFrame”的协议

Question

Hello SO and community!

Guess, my question somewhat resonates with this one .

However, trust the below task is a little bit different from that referenced above, namely to extract, transform, load data utilizing pandas.DataFrame , and I am stuck implementing Protocol for the purpose.

The code is below:

import io
import pandas as pd
import re
import requests
from functools import cache
from typing import Protocol
from zipfile import ZipFile
from pandas import DataFrame


@cache
def extract_can_from_url(url: str, **kwargs) -> DataFrame:
    '''
    Returns DataFrame from downloaded zip file from url
    Parameters
    ----------
    url : str
        url to download from.
    **kwargs : TYPE
        additional arguments to pass to pd.read_csv().
    Returns
    -------
    DataFrame
    '''
    name = url.split('/')[-1]
    if os.path.exists(name):
        with ZipFile(name, 'r').open(name.replace('-eng.zip', '.csv')) as f:
            return pd.read_csv(f, **kwargs)
    else:
        r = requests.get(url)
        with ZipFile(io.BytesIO(r.content)).open(name.replace('-eng.zip', '.csv')) as f:
            return pd.read_csv(f, **kwargs)


class ETL(Protocol):
    # =============================================================================
    # Maybe Using these items for dataclass:
    # url: str
    # meta: kwargs(default_factory=dict)
    # =============================================================================
    def __init__(self, url: str, **kwargs) -> None:
        return None

    def download(self) -> DataFrame:
        return DataFrame

    def retrieve_series_ids(self) -> list[str]:
        return list[str]

    def transform(self) -> DataFrame:
        return DataFrame

    def sum_up_series_ids(self) -> DataFrame:
        return DataFrame


class ETLCanadaFixedAssets(ETL):
    def __init__(self, url: str, **kwargs) -> None:
        self.url = url
        self.kwargs = kwargs

    @cache
    def download(self) -> DataFrame:
        self.df = extract_can_from_url(URL, index_col=0, usecols=range(14))
        return self.df

    def retrieve_series_ids(self) -> list[str]:
        # =========================================================================
        # Columns Specific to URL below, might be altered
        # =========================================================================
        self._columns = {
            "Prices": 0,
            "Industry": 1,
            "Flows and stocks": 2,
            "VECTOR": 3,
        }
        self.df_cut = self.df.loc[:, tuple(self._columns)]
        _q = (self.df_cut.iloc[:, 0].str.contains('2012 constant prices')) & \
            (self.df_cut.iloc[:, 1].str.contains('manufacturing', flags=re.IGNORECASE)) & \
            (self.df_cut.iloc[:, 2] == 'Linear end-year net stock')
        self.df_cut = self.df_cut[_q]
        self.series_ids = sorted(set(self.df_cut.iloc[:, -1]))
        return self.series_ids

    def transform(self) -> DataFrame:
        # =========================================================================
        # Columns Specific to URL below, might be altered
        # =========================================================================
        self._columns = {
            "VECTOR": 0,
            "VALUE": 1,
        }
        self.df = self.df.loc[:, tuple(self._columns)]
        self.df = self.df[self.df.iloc[:, 0].isin(self.series_ids)]
        return self.df

    def sum_up_series_ids(self) -> DataFrame:
        self.df = pd.concat(
            [
                self.df[self.df.iloc[:, 0] == series_id].iloc[:, [1]]
                for series_id in self.series_ids
            ],
            axis=1
        )
        self.df.columns = self.series_ids
        self.df['sum'] = self.df.sum(axis=1)
        return self.df.iloc[:, [-1]]

UPD

Instantiating the class ETLCanadaFixedAssets

df = ETLCanadaFixedAssets(URL, index_col=0, usecols=range(14)).download().retrieve_series_ids().transform().sum_up_series_ids()

returns an error, however, expected:

AttributeError: 'DataFrame' object has no attribute 'retrieve_series_ids'

Please can anyone provide a guidance for how to put these things together (namely how to retrieve the DataFrame which might have been retrieved otherwise using the procedural approach by calling the functions within the last class as they appear within the latter) and point at those mistakes which were made above?

Probably, there is another way to do this elegantly using injection.

Thank you very much in advance!

Answer 1

All the functions of ETLCanadaFixedAssets and ETL should return self so you can chain them together. You could add one more function that retrieves the encapsulated dataframe but that will always be called last, as the moment you call this function you cannot chain other functions any more. What you are trying to build is called fluent API you may read more about it here