enpi_api.l2.client.api.basket_api

  1import tempfile
  2from pathlib import Path
  3
  4import pandas as pd
  5from loguru import logger
  6
  7from enpi_api.l1 import openapi_client
  8from enpi_api.l2.client.api.file_api import FileApi
  9from enpi_api.l2.events.workflow_execution_task_waitable import WorkflowExecutionTaskWaitable
 10from enpi_api.l2.types.api_error import ApiErrorContext
 11from enpi_api.l2.types.basket import Basket, BasketExportFormat, BasketId, FastaExportConfig, TsvExportConfig
 12from enpi_api.l2.types.clone import CloneId
 13from enpi_api.l2.types.execution import Execution
 14from enpi_api.l2.types.log import LogLevel
 15from enpi_api.l2.types.sequence import SequenceId
 16from enpi_api.l2.types.task import TaskState
 17from enpi_api.l2.types.workflow import WorkflowExecutionId, WorkflowExecutionTaskId, WorkflowTaskTemplateName
 18
 19
 20class BasketApi:
 21    _inner_api_client: openapi_client.ApiClient
 22    _log_level: LogLevel
 23
 24    def __init__(self, inner_api_client: openapi_client.ApiClient, log_level: LogLevel):
 25        """@private"""
 26        self._inner_api_client = inner_api_client
 27        self._log_level = log_level
 28
 29    def get_baskets(self) -> list[Basket]:
 30        """Get all Baskets that belong to the user or were
 31        shared to them by other users.
 32
 33        Returns:
 34            list[enpi_api.l2.types.basket.Basket]: Available Baskets.
 35
 36        Raises:
 37            enpi_api.l2.types.api_error.ApiError: If API request fails.
 38
 39        Example:
 40            Get all available Baskets.
 41            ```python
 42            with EnpiApiClient() as enpi_client:
 43                baskets = enpi_client.basket_api.get_baskets()
 44                print(baskets)
 45            ```
 46        """
 47
 48        logger.info("Getting all available Baskets...")
 49
 50        basket_api_instance = openapi_client.BasketApi(self._inner_api_client)
 51
 52        with ApiErrorContext():
 53            get_baskets_response = basket_api_instance.get_baskets()
 54
 55        logger.success(f"Successfully got {len(get_baskets_response.baskets)} Baskets.")
 56
 57        return [Basket.from_raw(basket) for basket in get_baskets_response.baskets]
 58
 59    def get_basket(self, basket_id: BasketId) -> Basket:
 60        """Get a single Basket matched with the provided Basket ID.
 61
 62        Args:
 63            basket_id (enpi_api.l2.types.basket.BasketId): ID of a Basket to get.
 64
 65        Returns:
 66            enpi_api.l2.types.basket.Basket: Basket matching the provided ID.
 67
 68        Raises:
 69            enpi_api.l2.types.api_error.ApiError: If API request fails.
 70
 71        Example:
 72            Get a single Basket.
 73            ```python
 74            with EnpiApiClient() as enpi_client:
 75                id = BasketId(id)
 76                basket = enpi_client.basket_api.get_basket(id)
 77                print(basket)
 78            ```
 79        """
 80
 81        logger.info(f"Getting Basket with ID: '{basket_id}'...")
 82
 83        basket_api_instance = openapi_client.BasketApi(self._inner_api_client)
 84
 85        with ApiErrorContext():
 86            get_basket_response = basket_api_instance.get_basket(basket_id)
 87
 88        logger.success(f"Successfully got Basket with ID: '{basket_id}'.")
 89
 90        return Basket.from_raw(get_basket_response.basket)
 91
 92    def create_basket(self, name: str, shared: bool | None = True, description: str | None = "") -> Basket:
 93        """Create a new Basket.
 94
 95        Args:
 96            name (str): The name of the new Basket.
 97            shared (bool | None): Determines if the new Basket will be shared to other users in the organization.
 98                By default, all Baskets created via this function are being shared to other users.
 99            description (str | None): Description for the Basket.
100
101        Returns:
102            enpi_api.l2.types.basket.Basket: Newly created Basket.
103
104        Raises:
105            enpi_api.l2.types.api_error.ApiError: If API request fails.
106
107        Example:
108            Create a new Basket.
109            ```python
110            with EnpiApiClient() as enpi_client:
111                name = "New Basket"
112                basket = enpi_client.basket_api.create_basket(name)
113                print(basket)
114            ```
115        """
116
117        logger.info(f"Creating a new {'shared' if shared else 'private'} Basket with name: '{name}'...")
118
119        basket_api_instance = openapi_client.BasketApi(self._inner_api_client)
120
121        payload = openapi_client.CreateBasketRequestBody(name=name, shared=shared, description=description)
122
123        with ApiErrorContext():
124            create_basket_response = basket_api_instance.create_basket(payload)
125
126        logger.success(f"Successfully created a new Basket named: '{name}'.")
127
128        return Basket.from_raw(create_basket_response.basket)
129
130    def update_basket(
131        self,
132        basket_id: BasketId,
133        name: str | None = None,
134        shared: bool | None = None,
135        description: str | None = None,
136    ) -> None:
137        """Update properties of a Basket matched with the provided Basket ID.
138
139        Args:
140            basket_id (enpi_api.l2.types.basket.BasketId): ID of a Basket to be updated.
141            name (str | None): New name for the Basket. If no value is provided, Basket's name won't be changed.
142            shared (bool | None): Determines if a Basket will be shared to other users in the organization.
143                If no value is provided, Basket's shared status won't be changed.
144            description (str | None): New description for the Basket. If no value is provided, Basket's description won't be changed.
145
146        Raises:
147            enpi_api.l2.types.api_error.ApiError: If API request fails.
148
149        Example:
150            Update an existing Basket.
151            ```python
152            with EnpiApiClient() as enpi_client:
153                id = BasketId(id)
154                new_name = "Basket (renamed)"
155                is_basket_shared = True
156                enpi_client.basket_api.update_basket(id, new_name, is_basket_shared)
157            ```
158        """
159
160        logger.info(f"Updating Basket with ID: '{basket_id}'...")
161
162        basket_api_instance = openapi_client.BasketApi(self._inner_api_client)
163
164        payload = openapi_client.UpdateBasketRequestBody(name=name, shared=shared, description=description)
165
166        with ApiErrorContext():
167            basket_api_instance.update_basket(basket_id, payload)
168
169        logger.success(f"Successfully updated Basket with ID: '{basket_id}'.")
170
171    def delete_basket(self, basket_id: BasketId) -> None:
172        """Delete a Basket matched with the provided Basket ID.
173
174        Args:
175            basket_id (enpi_api.l2.types.basket.BasketId): ID of a Basket to be deleted.
176
177        Raises:
178            enpi_api.l2.types.api_error.ApiError: If API request fails.
179
180        Example:
181            Delete a Basket.
182            ```python
183            with EnpiApiClient() as enpi_client:
184                id = BasketId(id)
185                enpi_client.basket_api.delete_basket(id)
186            ```
187        """
188
189        logger.info(f"Deleting Basket with ID: '{basket_id}'...")
190
191        basket_api_instance = openapi_client.BasketApi(self._inner_api_client)
192
193        with ApiErrorContext():
194            basket_api_instance.delete_basket(basket_id)
195
196        logger.info(f"Successfully deleted Basket with ID:'{basket_id}'.")
197
198    def add_clones_to_basket(
199        self,
200        basket_id: BasketId,
201        clone_ids: list[CloneId] | None = None,
202        sequence_ids: list[SequenceId] | None = None,
203    ) -> Execution[None]:
204        """Resolve clones matched with provided clone and sequence IDs,
205        then add them into the target Basket.
206
207        > This functionality uses clone resolving.\n
208        > Clone resolving uses passed clone and sequence IDs in order to resolve clones.
209        > For each clone, a maximum of one *big* chain and one *small* chain sequence will be picked, resulting in a
210        maximum of two sequences per clone.
211        > Sequences matched with passed sequence IDs have priority over internally resolved sequences, meaning that if
212        possible, they will be picked as sequences for the resolved clones.
213
214        Args:
215            basket_id (enpi_api.l2.types.basket.BasketId): ID of a Basket to add clones to.
216            clone_ids (list[enpi_api.l2.types.clone.CloneId]): Clone IDs based on which clones will be resolved and added into the target Basket.
217            sequence_ids (list[enpi_api.l2.types.sequence.SequenceId]): Sequence IDs based on which clones will be resolved and added into
218            the target Basket. If clone resolving based on clone IDs and sequence IDs results in the same, "overlapping" clones (with the same clone IDs)
219            but potentially different sequences within, clones resolved with use of sequence IDs will be picked over the ones resolved with clone IDs.
220
221        Returns:
222            enpi_api.l2.types.execution.Execution[None]: An awaitable execution.
223
224        Raises:
225            ValueError: If the provided clone and/or sequence ID arrays are empty or invalid.
226            enpi_api.l2.types.api_error.ApiError: If API request fails.
227
228        Example:
229            Insert clones into a Basket.
230            ```python
231            with EnpiApiClient() as enpi_client:
232                basket_id = BasketId(1)
233                clone_ids = [
234                    CloneId(id) for id in [
235                        "d0c6982d-34cd-47cd-b465-a603dec0530c",
236                        "1c19fcbc-ec84-4baa-a641-42058df18303",
237                    ]
238                ]
239                sequence_ids = [SequenceId(id) for id in [100, 101, 102]]
240                enpi_client.basket_api.add_clones_to_basket(basket_id, clone_ids, sequence_ids).wait()
241            ```
242        """
243
244        logger.info(f"Adding clones to Basket with ID: '{basket_id}'...")
245
246        # Check if we got any ids to work with
247        if (clone_ids is None or len(clone_ids) == 0) and (sequence_ids is None or len(sequence_ids) == 0):
248            raise ValueError("Both clone and sequence IDs arrays are null, at least one of them needs to contain proper values.")
249
250        # Validate if ID types are right
251        if clone_ids is not None and not all([isinstance(id, str) for id in clone_ids]):
252            raise ValueError("Some of the passed clone IDs are not strings.")
253        elif sequence_ids is not None and not all([isinstance(id, int) for id in sequence_ids]):
254            raise ValueError("Some of the passed sequence IDs are not integers.")
255
256        basket_api_instance = openapi_client.BasketApi(self._inner_api_client)
257
258        payload = openapi_client.AddClonesRequestBody(
259            clone_ids=None if clone_ids is None else [str(id) for id in clone_ids],
260            sequence_ids=None if sequence_ids is None else [int(id) for id in sequence_ids],
261        )
262
263        with ApiErrorContext():
264            add_clones_response = basket_api_instance.add_clones(basket_id, payload)
265            assert add_clones_response.workflow_execution_id is not None
266
267            workflow_execution_id = WorkflowExecutionId(add_clones_response.workflow_execution_id)
268            waitable = WorkflowExecutionTaskWaitable[None](
269                workflow_execution_id=workflow_execution_id, task_template_name=WorkflowTaskTemplateName.ENPI_APP_BASKET_ADD_CLONES, on_complete=None
270            )
271
272        return Execution(wait=waitable.wait, check_execution_state=waitable.check_execution_state)
273
274    def remove_clones_from_basket(self, basket_id: BasketId, clone_ids: list[CloneId]) -> Execution[None]:
275        """Remove clones matches with provided clone IDs from the target Basket.
276
277        Args:
278            basket_id (enpi_api.l2.types.basket.BasketId): ID of a Basket from which clones will be removed.
279            clone_ids (list[enpi_api.l2.types.clone.CloneId]): IDs of clones that will be removed from target Basket.
280
281        Returns:
282            enpi_api.l2.types.execution.Execution[None]: An awaitable execution.
283
284        Raises:
285            ValueError: If the provided clone IDs array is invalid.
286            enpi_api.l2.types.api_error.ApiError: If API request fails.
287
288        Example:
289            Remove clones from a Basket.
290            ```python
291            with EnpiApiClient() as enpi_client:
292                basket_id = BasketId(1)
293                clone_ids = [
294                    CloneId(id) for id in [
295                        "d0c6982d-34cd-47cd-b465-a603dec0530c",
296                        "1c19fcbc-ec84-4baa-a641-42058df18303",
297                    ]
298                ]
299                enpi_client.basket_api.remove_clones_from_basket(basket_id, clone_ids).wait()
300            ```
301        """
302
303        n_of_clone_ids = len(clone_ids)
304        logger.info(f"Removing {n_of_clone_ids} clone{'' if n_of_clone_ids == 1 else 's'} from Basket with id: '{basket_id}'")
305
306        # Check if we got any ids to work with
307        if len(clone_ids) == 0:
308            raise ValueError("Provided clone IDs array is empty.")
309
310        # Validate if ID types are right
311        if not all([isinstance(id, str) for id in clone_ids]):
312            raise ValueError("Some of the passed clone IDs are not strings.")
313
314        basket_api_instance = openapi_client.BasketApi(self._inner_api_client)
315
316        payload = openapi_client.RemoveClonesRequestBody(
317            clone_ids=[str(id) for id in clone_ids],
318        )
319
320        with ApiErrorContext():
321            remove_clones_response = basket_api_instance.remove_clones(basket_id, payload)
322            assert remove_clones_response.workflow_execution_id is not None
323
324            workflow_execution_id = WorkflowExecutionId(remove_clones_response.workflow_execution_id)
325
326            waitable = WorkflowExecutionTaskWaitable[None](
327                workflow_execution_id=workflow_execution_id, task_template_name=WorkflowTaskTemplateName.ENPI_APP_BASKET_REMOVE_CLONES, on_complete=None
328            )
329
330        return Execution(wait=waitable.wait, check_execution_state=waitable.check_execution_state)
331
332    def export_basket_clones_as_tsv(
333        self,
334        basket_id: BasketId,
335        tsv_config: TsvExportConfig | None = None,
336        output_directory: Path | str | None = None,
337    ) -> Execution[Path]:
338        """Start a Basket clones export in a TSV format, download the result file and return a path to it.
339
340        Args:
341            basket_id (enpi_api.l2.types.basket.BasketId): ID of a Basket from which clones will be exported.
342            tsv_config (enpi_api.l2.types.basket.TsvExportConfig | None): The configuration for the TSV export.
343            output_directory (Path | str | None): Path to the directory in which the downloaded file will be stored.
344                If not provided, a temporary directory will be created.
345
346        Returns:
347            enpi_api.l2.types.execution.Execution[Path]: An awaitable execution that returns the local file path to the
348              exported file when awaited.
349
350        Raises:
351            enpi_api.l2.types.api_error.ApiError: If API request fails.
352
353        Example:
354            Export Basket clones into a TSV file.
355            ```python
356            with EnpiApiClient() as enpi_client:
357                id = BasketId(1)
358                export_dir = os.path.join(os.path.dirname(__file__), "example_dir")
359                path = enpi_client.basket_api.export_basket_clones_as_tsv(id, export_dir).wait()
360                print(path)
361            ```
362        """
363
364        logger.info(f"Exporting clones from Basket with ID: '{basket_id}' into a TSV file...")
365
366        basket_api_instance = openapi_client.BasketApi(self._inner_api_client)
367
368        payload = openapi_client.StartClonesExportRequestBody(
369            openapi_client.StartTsvClonesExportRequestBody(
370                format=BasketExportFormat.TSV, tsv_config=tsv_config.to_api_payload() if tsv_config is not None else None
371            )
372        )
373
374        with ApiErrorContext():
375            data = basket_api_instance.start_clones_export(basket_id, payload)
376            assert data.workflow_execution_id is not None
377
378            workflow_execution_id = WorkflowExecutionId(data.workflow_execution_id)
379
380            def on_complete(task_id: WorkflowExecutionTaskId, task_state: TaskState) -> Path:
381                with ApiErrorContext():
382                    file_api = FileApi(self._inner_api_client, self._log_level)
383                    file_path = file_api.download_export_by_workflow_execution_task_id(task_id=task_id, output_directory=output_directory)
384                    logger.success(f"Successfully exported clones from Basket with ID: '{basket_id}' into a TSV file: '{file_path}'.")
385                    return file_path
386
387            waitable = WorkflowExecutionTaskWaitable[Path](
388                workflow_execution_id=workflow_execution_id, on_complete=on_complete, task_template_name=WorkflowTaskTemplateName.ENPI_APP_BASKET_EXPORT
389            )
390
391            return Execution(wait=waitable.wait_and_return_result, check_execution_state=waitable.check_execution_state)
392
393    def export_basket_clones_as_df(self, basket_id: BasketId, tsv_config: TsvExportConfig | None = None) -> Execution[pd.DataFrame]:
394        """Start a Basket clones export in a TSV format, download the result and return
395            the data from it in a DataFrame object.
396
397        Args:
398            basket_id (enpi_api.l2.types.basket.BasketId): ID of a Basket from which clones will be exported.
399            tsv_config (enpi_api.l2.types.basket.TsvExportConfig | None): The configuration for the TSV export.
400
401        Returns:
402            enpi_api.l2.types.execution.Execution[pd.DataFrame]: An awaitable execution that returns a pandas DataFrame
403              object containing the exported data when awaited.
404
405        Raises:
406            enpi_api.l2.types.api_error.ApiError: If API request fails.
407
408        Example:
409            Export Basket clones into a DataFrame.
410            ```python
411            with EnpiApiClient() as client:
412                id = BasketId(1)
413                df = client.basket_api.export_basket_clones_as_df(id).wait()
414                print(df)
415            ```
416        """
417
418        logger.info(f"Exporting clones from Basket with ID: '{basket_id}' into a DataFrame...")
419
420        with tempfile.TemporaryDirectory() as temp_dir:
421            execution = self.export_basket_clones_as_tsv(basket_id=basket_id, tsv_config=tsv_config, output_directory=temp_dir)
422
423            def wait() -> pd.DataFrame:
424                tmp_file_path = execution.wait()
425                df = pd.read_csv(tmp_file_path, delimiter="\t")
426
427                logger.success(f"Successfully exported clones from Basket with ID: '{basket_id}' into a DataFrame.")
428
429                return df
430
431            return Execution(wait=wait, check_execution_state=execution.check_execution_state)
432
433    def export_basket_clones_as_fasta(
434        self,
435        basket_id: BasketId,
436        fasta_config: FastaExportConfig,
437        output_directory: Path | str | None = None,
438    ) -> Execution[Path]:
439        """Start a Basket clones export in a FASTA format, download the result and return the filepath to the file.
440
441        Args:
442            basket_id (enpi_api.l2.types.basket.BasketId): ID of a Basket from which clones will be exported.
443            fasta_config (enpi_api.l2.types.basket.FastaExportConfig): Configuration of the Basket FASTA export, determining the
444                shape of FASTA file headers and sequences. The enpi_api.l2.types.basket.fasta_config function is the prefered way of
445                a configuration object creation.
446            output_directory (Path | str | None): Path to a directory in which the downloaded file will be stored.
447                If not provided, a temporary directory will be created.
448
449        Returns:
450            enpi_api.l2.types.execution.Execution[Path]: An awaitable that returns information about clones added to the
451              Basket when awaited.
452
453        Raises:
454            enpi_api.l2.types.api_error.ApiError: If API request fails.
455
456        Example:
457            Export Basket clones into a FASTA file.
458            ```python
459            with EnpiApiClient() as client:
460                id = BasketId(1)
461                export_dir = os.path.join(os.path.dirname(__file__), "example_dir")
462                path = client.basket_api.export_basket_clones_as_fasta(
463                    basket_id,
464                    fasta_config(
465                        include_unique_clone_id_header=True,
466                        include_unique_sequence_id_header=True,
467                        include_chain_header=True,
468                        header_tag_keys=[
469                            "Full Sequence Amino Acids",
470                            "Organism",
471                        ],
472                        sequences=[
473                            FastaSequence(
474                                chain=Chain.HEAVY,
475                                tag_key=FastaExportSequenceKey.CDR3_AMINO_ACIDS,
476                            ),
477                            FastaSequence(
478                                chain=Chain.KAPPA,
479                                tag_key=FastaExportSequenceKey.CDR3_AMINO_ACIDS,
480                            ),
481                        ],
482                    ),
483                    output_directory=export_dir,
484                ).wait()
485                print(path)
486            ```
487        """
488
489        logger.info(f"Exporting clones from Basket with ID: '{basket_id}' into a FASTA file...")
490
491        basket_api_instance = openapi_client.BasketApi(self._inner_api_client)
492
493        payload = openapi_client.StartClonesExportRequestBody(
494            openapi_client.StartFastaClonesExportRequestBody(
495                format=BasketExportFormat.FASTA,
496                fasta_config=fasta_config.to_api_payload(),
497            )
498        )
499
500        with ApiErrorContext():
501            data = basket_api_instance.start_clones_export(basket_id, payload)
502            assert data.workflow_execution_id is not None
503
504            workflow_execution_id = WorkflowExecutionId(data.workflow_execution_id)
505
506            def on_complete(task_id: WorkflowExecutionTaskId, task_state: TaskState) -> Path:
507                file_api = FileApi(self._inner_api_client, self._log_level)
508                file_path = file_api.download_export_by_workflow_execution_task_id(task_id=task_id, output_directory=output_directory)
509
510                logger.success(f"Successfully exported clones from Basket with ID: '{basket_id}' into a FASTA file: '{file_path}'.")
511                return file_path
512
513            waitable = WorkflowExecutionTaskWaitable(
514                workflow_execution_id=workflow_execution_id, on_complete=on_complete, task_template_name=WorkflowTaskTemplateName.ENPI_APP_BASKET_EXPORT_FASTA
515            )
516
517            return Execution(wait=waitable.wait_and_return_result, check_execution_state=waitable.check_execution_state)