mardi 31 août 2021

How to avoid repeat code when converting files?

Let's say I have a class whose function is to convert file types:

@dataclass
class Converter:
    data: Union[str, pd.DataFrame]

    def to_pickle(self):
        """Check the type of data. 

        Returns:
            a pickle file if the table exists, a string otherwise.
        """
        if isinstance(self.data, pd.DataFrame):
            return pd.to_pickle(self.data, "table_data.pkl")
        else:
            return self.data
        
    def to_csv(self):
        """Check the type of data. 

        Returns:
            a csv file if the table exists, a string otherwise.
        """
        if isinstance(self.data, pd.DataFrame):
            return pd.to_csv(
                self.data, "./table_data.csv", 
                index=False, 
                encoding="utf_8_sig",)
        else:
            return self.data

Since both methods will first check the datatype. If the data is a dataframe, both methods will apply pd.to_csv and pd.to_pickle. Otherwise, a string will be returned.

It seems that the only difference between to_csv() and to_pickle() is the convert type (i.e. pd.to_csv and pd.to_pickle). Is there any way to not repeat the code?

Aucun commentaire:

Enregistrer un commentaire