cirro.sdk.file

  1import gzip
  2import json
  3import pickle
  4from io import BytesIO, StringIO
  5from pathlib import Path
  6from typing import List
  7
  8from typing import TYPE_CHECKING
  9if TYPE_CHECKING:
 10    import anndata
 11    from pandas import DataFrame
 12
 13from cirro.cirro_client import CirroApi
 14from cirro.models.file import File, PathLike
 15from cirro.sdk.asset import DataPortalAssets, DataPortalAsset
 16from cirro.sdk.exceptions import DataPortalInputError
 17from cirro.utils import convert_size
 18
 19
 20class DataPortalFile(DataPortalAsset):
 21    """
 22    Datasets are made up of a collection of File objects in the Data Portal.
 23    """
 24
 25    def __init__(self, file: File, client: CirroApi):
 26        """
 27        Instantiate by listing files from a dataset.
 28
 29        ```python
 30        from cirro import DataPortal
 31        portal = DataPortal()
 32        dataset = portal.get_dataset(
 33            project="id-or-name-of-project",
 34            dataset="id-or-name-of-dataset"
 35        )
 36        files = dataset.list_files()
 37        ```
 38        """
 39        # Attach the file object
 40        self._file = file
 41        self._client = client
 42
 43    # Note that the 'name' and 'id' attributes are set to the relative path
 44    # The purpose of this is to support the DataPortalAssets class functions
 45    @property
 46    def id(self) -> str:
 47        """Relative path of file within the dataset"""
 48        return self._file.relative_path
 49
 50    @property
 51    def name(self) -> str:
 52        """Relative path of file within the dataset"""
 53        return self._file.relative_path
 54
 55    @property
 56    def file_name(self) -> str:
 57        """Name of file, excluding the full folder path within the dataset"""
 58        return self._file.name
 59
 60    @property
 61    def relative_path(self) -> str:
 62        """Relative path of file within the dataset"""
 63        return self._file.relative_path
 64
 65    @property
 66    def absolute_path(self) -> str:
 67        """Fully URI to file object in AWS S3"""
 68        return self._file.absolute_path
 69
 70    @property
 71    def metadata(self) -> dict:
 72        """File metadata"""
 73        return self._file.metadata
 74
 75    @property
 76    def size_bytes(self) -> int:
 77        """File size (in bytes)"""
 78        return self._file.size
 79
 80    @property
 81    def size(self) -> str:
 82        """File size converted to human-readable (e.g., 4.50 GB)"""
 83        return convert_size(self._file.size)
 84
 85    def __str__(self):
 86        return f"{self.relative_path} ({self.size})"
 87
 88    def _get(self) -> bytes:
 89        """Internal method to call client.file.get_file"""
 90
 91        return self._client.file.get_file(self._file)
 92
 93    def read_csv(self, compression='infer', encoding='utf-8', **kwargs) -> 'DataFrame':
 94        """
 95        Parse the file as a Pandas DataFrame.
 96
 97        The default field separator is a comma (for CSV), use sep='\\t' for TSV.
 98
 99        File compression is inferred from the extension, but can be set
100        explicitly with the compression= flag.
101
102        All other keyword arguments are passed to pandas.read_csv
103        https://pandas.pydata.org/docs/reference/api/pandas.read_csv.html
104        """
105        import pandas
106
107        if compression == 'infer':
108            # If the file appears to be compressed
109            if self.relative_path.endswith('.gz'):
110                compression = dict(method='gzip')
111            elif self.relative_path.endswith('.bz2'):
112                compression = dict(method='bz2')
113            elif self.relative_path.endswith('.xz'):
114                compression = dict(method='xz')
115            elif self.relative_path.endswith('.zst'):
116                compression = dict(method='zstd')
117            else:
118                compression = None
119
120        if compression is not None:
121            handle = BytesIO(self._get())
122        else:
123            handle = StringIO(self._get().decode(encoding))
124
125        df = pandas.read_csv(
126            handle,
127            compression=compression,
128            encoding=encoding,
129            **kwargs
130        )
131        handle.close()
132        return df
133
134    def read_h5ad(self) -> 'anndata.AnnData':
135        """Read an AnnData object from a file."""
136        # Import the anndata library, and raise an error if it is not available
137        try:
138            import anndata as ad # noqa
139        except ImportError:
140            raise ImportError("The anndata library is required to read AnnData files. "
141                              "Please install it using 'pip install anndata'.")
142
143        # Download the file to a temporary file handle and parse the contents
144        with BytesIO(self._get()) as handle:
145            return ad.read_h5ad(handle)
146
147    def read_json(self, **kwargs):
148        """Read the file contents as a parsed JSON object (dict, list, etc.)."""
149        return json.loads(self._get(), **kwargs)
150
151    def read_parquet(self, **kwargs) -> 'DataFrame':
152        """
153        Read a Parquet file as a Pandas DataFrame.
154
155        Requires ``pyarrow`` or ``fastparquet`` to be installed.
156        All keyword arguments are passed to :func:`pandas.read_parquet`.
157        """
158        import pandas
159        return pandas.read_parquet(BytesIO(self._get()), **kwargs)
160
161    def read_feather(self, **kwargs) -> 'DataFrame':
162        """
163        Read a Feather file as a Pandas DataFrame.
164
165        Requires ``pyarrow`` to be installed.
166        All keyword arguments are passed to :func:`pandas.read_feather`.
167        """
168        import pandas
169        return pandas.read_feather(BytesIO(self._get()), **kwargs)
170
171    def read_pickle(self, **kwargs):
172        """Read the file contents as a Python pickle object."""
173        return pickle.loads(self._get(), **kwargs)
174
175    def read_excel(self, **kwargs) -> 'DataFrame':
176        """
177        Read an Excel file (``.xlsx`` / ``.xls``) as a Pandas DataFrame.
178
179        Requires ``openpyxl`` (for ``.xlsx``) or ``xlrd`` (for ``.xls``).
180        All keyword arguments are passed to :func:`pandas.read_excel`.
181        """
182        import pandas
183        return pandas.read_excel(BytesIO(self._get()), **kwargs)
184
185    def readlines(self, encoding='utf-8', compression=None) -> List[str]:
186        """Read the file contents as a list of lines."""
187
188        return self.read(
189            encoding=encoding,
190            compression=compression
191        ).splitlines()
192
193    def read(self, encoding='utf-8', compression=None) -> str:
194        """Read the file contents as text."""
195
196        # Get the raw file contents
197        cont = self._get()
198
199        # If the file is uncompressed
200        if compression is None:
201            return cont.decode(encoding)
202        # If the file is compressed
203        else:
204
205            # Only gzip-compression is supported currently
206            if compression != "gzip":
207                raise DataPortalInputError("compression may be 'gzip' or None")
208
209            with gzip.open(
210                BytesIO(
211                    cont
212                ),
213                'rt',
214                encoding=encoding
215            ) as handle:
216                return handle.read()
217
218    def read_bytes(self) -> BytesIO:
219        """Get a generic BytesIO object representing the Data Portal File, to be passed into readers."""
220        return BytesIO(self._get())
221
222    def download(self, download_location: str = None) -> Path:
223        """
224        Download the file to a local directory.
225
226        Returns:
227            Path to download file
228        """
229
230        if download_location is None:
231            raise DataPortalInputError("Must provide download location")
232
233        return self._client.file.download_files(
234            self._file.access_context,
235            download_location,
236            [self.relative_path]
237        )[0]
238
239    def validate(self, local_path: PathLike):
240        """
241        Validate that the local file matches the remote file by comparing checksums.
242
243        Args:
244            local_path (PathLike): Path to the local file to validate
245        Raises:
246            ValueError: If checksums do not match
247            RuntimeWarning: If the remote checksum is not available or not supported
248        """
249        self._client.file.validate_file(self._file, local_path)
250
251    def is_valid(self, local_path: PathLike) -> bool:
252        """
253        Check if the local file matches the remote file by comparing checksums.
254
255        Args:
256            local_path (PathLike): Path to the local file to validate
257        Returns:
258            bool: True if the local file matches the remote file, False otherwise
259        Raises:
260            RuntimeWarning: If the remote checksum is not available or not supported
261        """
262        if not local_path:
263            raise DataPortalInputError("Must provide local path to validate file")
264
265        return self._client.file.is_valid_file(self._file, local_path)
266
267
268class DataPortalFiles(DataPortalAssets[DataPortalFile]):
269    """Collection of DataPortalFile objects."""
270
271    asset_name = "file"
272
273    def download(self, download_location: str = None) -> List[Path]:
274        """
275        Download the collection of files to a local directory.
276
277        Returns:
278            List of paths to downloaded files.
279        """
280
281        local_paths = []
282        for f in self:
283            local_paths.append(f.download(download_location))
284        return local_paths