pybiolib 0.2.951__py3-none-any.whl → 1.2.1890__py3-none-any.whl

biolib/utils/cache_state.py

@@ -0,0 +1,94 @@
+import abc
+import json
+import os
+import time
+from datetime import datetime, timezone
+
+import appdirs  # type: ignore
+
+from biolib.biolib_errors import BioLibError
+from biolib.biolib_logging import logger_no_user_data
+from biolib.typing_utils import Generic, Optional, TypeVar
+
+StateType = TypeVar('StateType')  # pylint: disable=invalid-name
+
+
+class CacheStateError(BioLibError):
+    pass
+
+
+class CacheState(abc.ABC, Generic[StateType]):
+    @property
+    @abc.abstractmethod
+    def _state_path(self) -> str:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def _get_default_state(self) -> StateType:
+        raise NotImplementedError
+
+    @property
+    def _user_cache_dir(self) -> str:
+        user_cache_dir: str = appdirs.user_cache_dir(appname='pybiolib', appauthor='biolib')
+        os.makedirs(user_cache_dir, exist_ok=True)
+        return user_cache_dir
+
+    @property
+    def _state_lock_path(self) -> str:
+        return f'{self._state_path}.lock'
+
+    def __init__(self, fail_fast_on_lock_acquire: bool = False) -> None:
+        self._state: Optional[StateType] = None
+        self._fail_fast_on_lock_acquire: bool = fail_fast_on_lock_acquire
+
+    def __enter__(self) -> StateType:
+        logger_no_user_data.debug(f'CacheState: Entering state path: {self._state_path}...')
+        try:
+            self._acquire_state_lock()
+            if os.path.exists(self._state_path):
+                with open(self._state_path, mode='r') as file:
+                    self._state = json.loads(file.read())
+            else:
+                self._state = self._get_default_state()
+                with open(self._state_path, mode='w') as file:
+                    file.write(json.dumps(self._state))
+
+            # Check for type checking
+            if self._state is None:
+                raise CacheStateError('Internal state is not defined')
+        except BaseException as error:  # pylint: disable=broad-except
+            logger_no_user_data.debug(f'Could not get LFS lock, got error: {error}...')
+            raise error
+        return self._state
+
+    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
+        with open(self._state_path, mode='w') as file:
+            file.write(json.dumps(self._state))
+
+        self._release_state_lock()
+        logger_no_user_data.debug(f'CacheState: Exited state path: {self._state_path}')
+
+    def _acquire_state_lock(self) -> None:
+        for _ in range(10):
+            try:
+                lock_file = open(self._state_lock_path, mode='x')
+                lock_file.close()
+                return
+            except BaseException as error:  # pylint: disable=broad-except
+                logger_no_user_data.debug(f'Failed to acquire lock file "{self._state_lock_path}". Got error: {error}')
+                if self._fail_fast_on_lock_acquire:
+                    raise CacheStateError(f'Failed to acquire lock file "{self._state_lock_path}": {error}') from error
+
+            time.sleep(0.5)
+
+        raise CacheStateError(f'Cache state timed out waiting to acquire lock file "{self._state_lock_path}"')
+
+    def _release_state_lock(self) -> None:
+        if os.path.exists(self._state_lock_path):
+            os.remove(self._state_lock_path)
+        else:
+            raise CacheStateError('Cache state was not locked.')
+
+    @staticmethod
+    def get_timestamp_now() -> str:
+        return datetime.now(timezone.utc).isoformat()

biolib/utils/multipart_uploader.py

@@ -0,0 +1,194 @@
+import math
+import multiprocessing
+import multiprocessing.pool
+import os
+import time
+from urllib.parse import urlparse
+
+import biolib.api
+from biolib._internal.http_client import HttpClient
+from biolib.biolib_api_client import BiolibApiClient
+from biolib.biolib_errors import BioLibError
+from biolib.biolib_logging import logger, logger_no_user_data
+from biolib.typing_utils import Callable, Dict, Iterator, List, Optional, Tuple, TypedDict
+
+
+def get_chunk_iterator_from_bytes(byte_buffer: bytes, chunk_size_in_bytes: int = 50_000_000) -> Iterator[bytes]:
+    chunk_count = math.ceil(len(byte_buffer) / chunk_size_in_bytes)
+    for chunk_number in range(chunk_count):
+        start = chunk_size_in_bytes * chunk_number
+        stop = start + chunk_size_in_bytes
+        yield byte_buffer[start:stop]
+
+
+def get_chunk_iterator_from_file_object(file_object, chunk_size_in_bytes: int = 50_000_000) -> Iterator[bytes]:
+    while True:
+        data = file_object.read(chunk_size_in_bytes)
+        if not data:
+            break
+        yield data
+
+
+class RequestOptions(TypedDict):
+    headers: Optional[Dict[str, str]]
+    requires_biolib_auth: bool
+    path: str
+
+
+class _PartMetadata(TypedDict):
+    ETag: str
+    PartNumber: int
+
+
+_UploadChunkInputType = Tuple[int, bytes]
+_UploadChunkReturnType = Tuple[_PartMetadata, int]
+
+
+class MultiPartUploader:
+    def __init__(
+        self,
+        complete_upload_request: RequestOptions,
+        get_presigned_upload_url_request: RequestOptions,
+        start_multipart_upload_request: Optional[RequestOptions] = None,
+        use_process_pool: Optional[bool] = None,
+        on_progress: Optional[Callable[[int, int], None]] = None,
+    ):
+        self._complete_upload_request = complete_upload_request
+        self._get_presigned_upload_url_request = get_presigned_upload_url_request
+        self._start_multipart_upload_request = start_multipart_upload_request
+        self._bytes_uploaded: int = 0
+        self._use_process_pool = use_process_pool
+        self._on_progress = on_progress
+
+    def upload(self, payload_iterator: Iterator[bytes], payload_size_in_bytes: int) -> None:
+        parts: List[_PartMetadata] = []
+
+        iterator_with_index: Iterator[_UploadChunkInputType] = enumerate(payload_iterator, 1)  # type: ignore
+        logger_no_user_data.debug(f'Starting multipart upload of payload with size {payload_size_in_bytes} bytes')
+
+        if self._start_multipart_upload_request:
+            try:
+                biolib.api.client.post(
+                    authenticate=self._start_multipart_upload_request['requires_biolib_auth'],
+                    headers=self._start_multipart_upload_request['headers'],
+                    path=self._start_multipart_upload_request['path'],
+                )
+            except BaseException as error:
+                logger_no_user_data.debug(f'Failed to start multipart upload got error: {error}')
+                raise error
+
+        # if multiprocessing start method is spawn or we are running in a daemon process,
+        # multiprocessing.Pool may fail when called from script
+        if multiprocessing.get_start_method() == 'spawn' or multiprocessing.current_process().daemon:
+            logger_no_user_data.debug('Uploading multipart from main process...')
+            for chunk in iterator_with_index:
+                upload_chunk_response = self._upload_chunk(chunk)
+                self._update_progress_bar_and_parts(
+                    upload_chunk_response=upload_chunk_response,
+                    parts=parts,
+                    payload_size_in_bytes=payload_size_in_bytes,
+                )
+        else:
+            # use 16 cores, unless less is available
+            pool_size = min(16, multiprocessing.cpu_count() - 1)
+            process_pool = (
+                multiprocessing.Pool(pool_size)
+                if self._use_process_pool
+                else multiprocessing.pool.ThreadPool(pool_size)
+            )
+
+            try:
+                response: _UploadChunkReturnType
+                for response in process_pool.imap(self._upload_chunk, iterator_with_index):
+                    self._update_progress_bar_and_parts(
+                        upload_chunk_response=response, parts=parts, payload_size_in_bytes=payload_size_in_bytes
+                    )
+            finally:
+                logger_no_user_data.debug('Multipart upload closing process pool...')
+                process_pool.close()
+
+        requires_biolib_auth = self._complete_upload_request['requires_biolib_auth']
+        if requires_biolib_auth:
+            BiolibApiClient.refresh_auth_token()
+
+        logger_no_user_data.debug(f'Uploaded {len(parts)} parts, now calling complete upload...')
+        biolib.api.client.post(
+            authenticate=requires_biolib_auth,
+            headers=self._complete_upload_request['headers'],
+            data={'parts': parts, 'size_bytes': self._bytes_uploaded},
+            path=self._complete_upload_request['path'],
+        )
+
+    def _upload_chunk(self, _input: _UploadChunkInputType) -> _UploadChunkReturnType:
+        part_number, chunk = _input
+        requires_biolib_auth = self._get_presigned_upload_url_request['requires_biolib_auth']
+
+        for index in range(20):  # will fail after approximately sum_i(i^2+2) = 41 min if range (20)
+            if requires_biolib_auth:
+                BiolibApiClient.refresh_auth_token()
+
+            logger_no_user_data.debug(f'Uploading part number {part_number} with size {len(chunk)} bytes...')
+            presigned_upload_url = None
+            try:
+                logger_no_user_data.debug(f'Getting upload URL for chunk {part_number}...')
+                get_url_response = biolib.api.client.get(
+                    authenticate=requires_biolib_auth,
+                    headers=self._get_presigned_upload_url_request['headers'],
+                    params={'part_number': part_number},
+                    path=self._get_presigned_upload_url_request['path'],
+                )
+
+                presigned_upload_url = get_url_response.json()['presigned_upload_url']
+
+            except Exception as error:  # pylint: disable=broad-except
+                logger_no_user_data.warning(f'Error when getting url for part {part_number}. Retrying...')
+                logger.debug(f'Upload error: {error}')
+
+            if presigned_upload_url:
+                try:
+                    app_caller_proxy_job_storage_base_url = os.getenv('BIOLIB_CLOUD_JOB_STORAGE_BASE_URL', '')
+                    if app_caller_proxy_job_storage_base_url:
+                        # Done to hit App Caller Proxy when uploading result from inside an app
+                        parsed_url = urlparse(presigned_upload_url)
+                        presigned_upload_url = (
+                            f'{app_caller_proxy_job_storage_base_url}{parsed_url.path}?{parsed_url.query}'
+                        )
+
+                    put_chunk_response = HttpClient.request(
+                        url=presigned_upload_url,
+                        data=chunk,
+                        method='PUT',
+                        timeout_in_seconds=300,
+                    )
+                    return _PartMetadata(PartNumber=part_number, ETag=put_chunk_response.headers['ETag']), len(chunk)
+
+                except Exception as error:  # pylint: disable=broad-except
+                    logger_no_user_data.warning(f'Encountered error when uploading part {part_number}. Retrying...')
+                    logger.debug(f'Upload error: {error} ({presigned_upload_url})')
+
+            time.sleep(index * index + 2)
+
+        logger_no_user_data.debug(f'Max retries hit, when uploading part {part_number}. Exiting...')
+        raise BioLibError(f'Max retries hit, when uploading part {part_number}. Exiting...')
+
+    def _update_progress_bar_and_parts(
+        self,
+        upload_chunk_response: _UploadChunkReturnType,
+        parts: List[_PartMetadata],
+        payload_size_in_bytes: int,
+    ) -> None:
+        part_metadata, chunk_byte_length = upload_chunk_response
+        part_number = part_metadata['PartNumber']
+
+        parts.append(part_metadata)
+        self._bytes_uploaded += chunk_byte_length
+
+        if self._on_progress is not None:
+            self._on_progress(self._bytes_uploaded, payload_size_in_bytes)
+
+        approx_progress_percent = min(self._bytes_uploaded / (payload_size_in_bytes + 1) * 100, 100)
+        approx_rounded_progress = round(approx_progress_percent, 2)
+        logger_no_user_data.debug(
+            f'Uploaded part number {part_number} with size {chunk_byte_length} bytes, '
+            f'the approximate progress is {approx_rounded_progress}%'
+        )

biolib/utils/seq_util.py

@@ -0,0 +1,150 @@
+import re
+from io import BufferedIOBase, TextIOBase
+
+from biolib.typing_utils import Dict, Iterator, List, Optional, Union
+
+
+class SeqUtilRecord:
+    def __init__(
+        self,
+        sequence: str,
+        sequence_id: str,
+        description: Optional['str'] = None,
+        properties: Optional[Dict[str, str]] = None,
+    ):
+        self.sequence = sequence
+        self.id = sequence_id  # pylint: disable=invalid-name
+        self.description = description
+
+        if properties:
+            disallowed_pattern = re.compile(r'[=\[\]\n]')
+            for key, value in properties.items():
+                assert not bool(disallowed_pattern.search(key)), 'Key cannot contain characters =[] and newline'
+                assert not bool(disallowed_pattern.search(value)), 'Value cannot contain characters =[] and newline'
+            self.properties = properties
+        else:
+            self.properties = {}
+
+    def __repr__(self) -> str:
+        return f'{self.__class__.__name__} ({self.id})'
+
+
+class SeqUtil:
+    @staticmethod
+    def parse_fasta(
+        input_file: Union[str, BufferedIOBase, None] = None,
+        default_header: Optional[str] = None,
+        allow_any_sequence_characters: bool = False,
+        use_strict_alphabet: Optional[bool] = False,
+        allow_empty_sequence: bool = True,
+        file_name: Optional[str] = None,
+) -> Iterator[SeqUtilRecord]:
+        def process_and_yield_record(header: str, sequence_lines: List[str]):
+            sequence = ''.join(sequence_lines)
+            sequence_id = header.split()[0]
+            if allow_any_sequence_characters and use_strict_alphabet:
+                raise Exception(
+                    'Error: Please choose either allow_any_sequence_characters or use_strict_alphabet'
+                )
+            if not allow_any_sequence_characters:
+                if use_strict_alphabet:
+                    invalid_sequence_characters = SeqUtil._find_invalid_sequence_characters_strict(sequence)
+                else:
+                    invalid_sequence_characters = SeqUtil._find_invalid_sequence_characters(sequence)
+                if invalid_sequence_characters:
+                    raise Exception(
+                        f'Error: Invalid character ("{invalid_sequence_characters[0]}") found in sequence {sequence_id}'
+                    )
+            if not allow_empty_sequence and not sequence:
+                raise Exception(f'Error: No sequence found for fasta entry {sequence_id}')
+            yield SeqUtilRecord(
+                sequence=sequence,
+                sequence_id=sequence_id,
+                description=header[len(sequence_id):].strip()
+            )
+
+        def line_generator_from_buffered_io_base(file_handle: BufferedIOBase) -> Iterator[str]:
+            for line in file_handle:
+                yield line.decode('utf-8')
+
+        def line_generator_from_text_io_base(file_handle: TextIOBase) -> Iterator[str]:
+            for line in file_handle:
+                yield line
+
+        if input_file is None:
+            if file_name:
+                input_file = file_name
+            else:
+                raise ValueError('input_file must be a file name (str) or a BufferedIOBase object')
+
+        file_handle = None
+        if isinstance(input_file, str):
+            file_handle = open(input_file, "rb")
+            line_iterator = line_generator_from_buffered_io_base(file_handle)
+        elif isinstance(input_file, BufferedIOBase):
+            line_iterator = line_generator_from_buffered_io_base(input_file)
+        elif isinstance(input_file, TextIOBase):
+            line_iterator = line_generator_from_text_io_base(input_file)
+        else:
+            raise ValueError('input_file must be a file name (str) or a BufferedIOBase object')
+
+        header = None
+        sequence_lines: List[str] = []
+
+        try:
+            for line_number, line in enumerate(line_iterator):
+                line = line.strip()
+                if not line:
+                    continue # skip empty lines
+                if line.startswith('>'):
+                    if header is not None:
+                        yield from process_and_yield_record(header, sequence_lines)
+
+                    header = line[1:].strip()
+                    sequence_lines = []
+                else:
+                    if header is None:
+                        if default_header:
+                            yield from process_and_yield_record(f"{default_header}{line_number}", [line])
+                        else:
+                            raise Exception(f'No header line found in FASTA file "{file_name}"')
+                    else:
+                        sequence_lines.append(line)
+
+            if header is not None:
+                yield from process_and_yield_record(header, sequence_lines)
+        finally:
+            if file_handle:
+                file_handle.close()
+
+    @staticmethod
+    def write_records_to_fasta(file_name: str, records: List[SeqUtilRecord]) -> None:
+        with open(file_name, mode='w') as file_handle:
+            for record in records:
+                optional_description = f' {record.description}' if record.description else ''
+                if record.properties:
+                    for key, value in record.properties.items():
+                        optional_description += f' [{key}={value}]'
+                sequence = '\n'.join(record.sequence[i : i + 80] for i in range(0, len(record.sequence), 80))
+                file_handle.write(f'>{record.id}{optional_description}\n{sequence}\n')
+
+    @staticmethod
+    def _find_invalid_sequence_characters(sequence: str) -> List[str]:
+        allowed_sequence_chars = set('abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789-_.')
+        invalid_chars = [char for char in sequence if char not in allowed_sequence_chars]
+        return invalid_chars
+
+    @staticmethod
+    def _find_invalid_sequence_characters_strict(sequence: str) -> List[str]:
+        # Equivalent to fair-esm alphabet, compatible with ESM-models
+        # Excludes digits, '_' and 'J' (ambiguous letter only used in mass-spec NMR)
+        # https://github.com/facebookresearch/esm/blob/2b369911bb5b4b0dda914521b9475cad1656b2ac/esm/constants.py#L8
+        allowed_sequence_chars = set('lagvsertidpkqnfymhwcxbuzoLAGVSERTIDPKQNFYMHWCXBUZO-.')
+        invalid_chars = [char for char in sequence if char not in allowed_sequence_chars]
+        return invalid_chars
+
+    @staticmethod
+    def _find_invalid_sequence_id_characters(sequence: str) -> List[str]:
+        allowed_chars = set('abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789-_.:*#')
+        invalid_chars = [char for char in sequence if char not in allowed_chars]
+        return invalid_chars