my_project.fileManager

  1import pandas as pd
  2
  3class FileManager():
  4    """
  5    A simple file manager.
  6        
  7    Parameters
  8    ----------
  9    format_ : str, optional
 10        Initial format of the files its managing
 11    
 12    Attributes
 13    ----------
 14    format : str | None
 15        Format of the files its managing
 16        
 17    Examples
 18    --------
 19    >>> fm = FileManager()
 20    >>> fm.set_format("parquet")
 21    >>> df = fm.read("parquet_data")
 22    
 23    >>> fm = FileManager("csv")
 24    >>> df = fm.read("csv_data")
 25    """
 26        
 27    formats = {
 28        'csv': {'write': lambda df, f: df.to_csv(f, index=False), 'read': pd.read_csv, 'ext': '.csv'},
 29        'parquet': {'write': lambda df, f: df.to_parquet(f), 'read': pd.read_parquet, 'ext': '.parquet'},
 30        'feather': {'write': lambda df, f: df.to_feather(f), 'read': pd.read_feather, 'ext': '.feather'},
 31        'pickle': {'write': lambda df, f: df.to_pickle(f), 'read': pd.read_pickle, 'ext': '.pkl'}
 32    }
 33
 34    def __init__(self,*, format_: str | None = None):
 35        self.format = None
 36        if format_ is not None:
 37            self.set_format(format_)
 38        
 39    @staticmethod
 40    def avalible_formats() -> list[str]:
 41        """
 42        Static method. Obtains the supported file formats.
 43        
 44        Returns
 45        -------
 46        list[str]
 47            List of strings, each a supported format
 48        
 49        Examples
 50        --------
 51        >>> formats = FileManager.avalible_formats()
 52        """
 53        
 54        return list(FileManager.formats.keys())
 55        
 56    @staticmethod
 57    def is_format(format_: str) -> bool:
 58        """
 59        Static method. Checks if a file format is supported.
 60        
 61        Parameters
 62        ----------
 63        format_ : str
 64            File format to check
 65        
 66        Returns
 67        -------
 68        bool
 69            True if it is supported
 70        
 71        Examples
 72        --------
 73        >>> csv_supported = FileManager.is_format("csv")
 74        """
 75        
 76        return (format_ in FileManager.formats.keys())
 77
 78    def set_format(self, format_: str) -> None:
 79        """
 80        Sets the file format of the files to manage.
 81        
 82        Parameters
 83        ----------
 84        format_ : str
 85            File format to use
 86        
 87        Examples
 88        --------
 89        >>> fm.set_format("csv")
 90
 91        Raises
 92        ------
 93        Exception
 94            If the format is not supported
 95        """
 96        
 97        if not self.is_format(format_):
 98            raise f"{format_} format not supported. Avalible formats: {self.avalible_formats()}"
 99        self.format = format_
100
101    def _get_name(self, input_filename: str) -> str:
102        """
103        Adds the apropiate extension to the filename.
104        
105        Parameters
106        ----------
107        input_filename : str
108            Filename without extension
109        
110        Returns
111        -------
112        str
113            Filename with the extension of the current format
114        
115        Examples
116        --------
117        >>> fm.set_format("csv")
118        >>> csv_filename = fm._get_name("name")
119        """
120        
121        filename = f'{input_filename}{self.formats[self.format]["ext"]}'
122        return filename
123        
124    def read(self, input_filename: str) -> pd.DataFrame:
125        """
126        Reads from the selected file.
127        
128        Parameters
129        ----------
130        input_filename : str
131            Filename without extension
132        
133        Returns
134        -------
135        pd.DataFrame
136            Dataframe with the content read from the selected file
137        
138        Examples
139        --------
140        >>> fm.set_format("csv")
141        >>> df = fm.read("csv_file")
142        """
143        
144        filename = self._get_name(input_filename)
145        return self.formats[self.format]["read"](filename)
146        
147    def write(self, df: pd.DataFrame, input_filename: str) -> str | None:
148        """
149        Writes to the selected file.
150        
151        Parameters
152        ----------
153        df : pd.DataFrame
154            Dataframe containg the content to write
155        input_filename : str
156            Filename without extension
157        
158        Returns
159        -------
160        str | None
161            Passes the return of the apropiate write call
162        
163        Examples
164        --------
165        >>> fm.set_format("csv")
166        >>> fm.write(df, "csv_file")
167        """
168        
169        filename = self._get_name(input_filename)
170        return self.formats[self.format]["write"](df, filename)