pybiolib 1.2.166__tar.gz → 1.2.174__tar.gz

{pybiolib-1.2.166 → pybiolib-1.2.174}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: pybiolib
-Version: 1.2.166
+Version: 1.2.174
 Summary: BioLib Python Client
 Home-page: https://github.com/biolib
 License: MIT

{pybiolib-1.2.166 → pybiolib-1.2.174}/biolib/_data_record/data_record.py

@@ -1,10 +1,9 @@
 import os
 from collections import namedtuple
 from datetime import datetime
-from fnmatch import fnmatch
 from pathlib import Path
 from struct import Struct
-from typing import Callable, Dict, List, Optional, Union, cast
+from typing import Callable, Dict, Iterable, List, Optional, Union, cast
 
 from biolib import api
 from biolib._internal import types
@@ -16,6 +15,7 @@ from biolib._internal.data_record.push_data import (
 )
 from biolib._internal.data_record.remote_storage_endpoint import DataRecordRemoteStorageEndpoint
 from biolib._internal.http_client import HttpClient
+from biolib._internal.types.file_node import ZipFileNodeDict
 from biolib.api import client as api_client
 from biolib.biolib_api_client import BiolibApiClient
 from biolib.biolib_api_client.lfs_types import DataRecordInfo, DataRecordVersion, DataRecordVersionInfo
@@ -23,9 +23,8 @@ from biolib.biolib_binary_format import LazyLoadedFile
 from biolib.biolib_binary_format.utils import RemoteIndexableBuffer
 from biolib.biolib_logging import logger
 from biolib.utils.app_uri import parse_app_uri
-from biolib.utils.zip.remote_zip import RemoteZip
 
-PathFilter = Union[str, Callable[[str], bool]]
+PathFilter = Union[str, List[str], Callable[[str], bool]]
 
 
 class DataRecord:
@@ -51,17 +50,25 @@ class DataRecord:
 
         return uri_parsed['app_name']
 
-    def list_files(self, path_filter: Optional[PathFilter] = None) -> List[LazyLoadedFile]:
-        remote_storage_endpoint = DataRecordRemoteStorageEndpoint(
-            resource_version_uuid=self._state['resource_version_uuid'],
+    def list_files(
+        self,
+        path_filter: Optional[PathFilter] = None,
+        max_count: Optional[int] = 100_000,
+    ) -> List[LazyLoadedFile]:
+        files = list(
+            self._fetch_files(
+                path_filter=path_filter,
+                max_count=max_count + 1 if max_count is not None else None,
+            )
         )
-        files: List[LazyLoadedFile] = []
-        with RemoteZip(url=remote_storage_endpoint.get_remote_url()) as remote_zip:
-            central_directory = remote_zip.get_central_directory()
-            for file_info in central_directory.values():
-                files.append(self._get_file(remote_storage_endpoint, file_info))
 
-        return self._get_filtered_files(files=files, path_filter=path_filter) if path_filter else files
+        if max_count is not None and len(files) > max_count:
+            raise Exception(
+                f'list_files returned more than {max_count} files. '
+                f'Please set the keyword argument "max_count" to a higher number.'
+            )
+
+        return files
 
     def download_zip(self, output_path: str):
         remote_storage_endpoint = DataRecordRemoteStorageEndpoint(
@@ -188,8 +195,50 @@ class DataRecord:
             for result in results
         ]
 
+    def _fetch_files(
+        self,
+        max_count: Optional[int],
+        path_filter: Optional[PathFilter] = None,
+    ) -> Iterable[LazyLoadedFile]:
+        if path_filter and not (isinstance(path_filter, (str, list)) or callable(path_filter)):
+            raise Exception('Expected path_filter to be a string, a list of strings or a function')
+
+        path_filters = (
+            [path_filter] if isinstance(path_filter, str) else path_filter if isinstance(path_filter, list) else []
+        )
+
+        resource_version_uuid = self._state['resource_version_uuid']
+        remote_storage_endpoint = DataRecordRemoteStorageEndpoint(resource_version_uuid)
+
+        page: Optional[int] = 1
+        yielded_files: int = 0
+        while page:
+            response = api.client.post(
+                path=f'/proxy/files/data-record-versions/{resource_version_uuid}/query/',
+                data=dict(page=page, page_size=1_000, path_filters=path_filters),
+            ).json()
+
+            for file_node_dict in cast(List[ZipFileNodeDict], response['results']):
+                if file_node_dict['is_dir']:
+                    continue
+
+                if callable(path_filter) and not path_filter(file_node_dict['dir_path'] + file_node_dict['name']):
+                    continue
+
+                yield self._get_file(remote_storage_endpoint, file_node_dict)
+                yielded_files += 1
+
+                if max_count is not None and yielded_files >= max_count:
+                    page = None
+                    break
+
+            page = page + 1 if page is not None and response['page_count'] > page else None
+
     @staticmethod
-    def _get_file(remote_storage_endpoint: DataRecordRemoteStorageEndpoint, file_info: Dict) -> LazyLoadedFile:
+    def _get_file(
+        remote_storage_endpoint: DataRecordRemoteStorageEndpoint,
+        file_node_dict: ZipFileNodeDict,
+    ) -> LazyLoadedFile:
         local_file_header_signature_bytes = b'\x50\x4b\x03\x04'
         local_file_header_struct = Struct('<H2sHHHIIIHH')
         LocalFileHeader = namedtuple(
@@ -208,7 +257,7 @@ class DataRecord:
             ),
         )
 
-        local_file_header_start = file_info['header_offset'] + len(local_file_header_signature_bytes)
+        local_file_header_start = file_node_dict['zip_meta']['header_start'] + len(local_file_header_signature_bytes)
         local_file_header_end = local_file_header_start + local_file_header_struct.size
 
         def file_start_func() -> int:
@@ -227,26 +276,11 @@ class DataRecord:
 
         return LazyLoadedFile(
             buffer=RemoteIndexableBuffer(endpoint=remote_storage_endpoint),
-            length=file_info['file_size'],
-            path=file_info['filename'],
+            length=file_node_dict['zip_meta']['size_on_disk'],
+            path=file_node_dict['dir_path'] + file_node_dict['name'],
             start=None,
             start_func=file_start_func,
         )
 
-    @staticmethod
-    def _get_filtered_files(files: List[LazyLoadedFile], path_filter: PathFilter) -> List[LazyLoadedFile]:
-        if not (isinstance(path_filter, str) or callable(path_filter)):
-            raise Exception('Expected path_filter to be a string or a function')
-
-        if callable(path_filter):
-            return list(filter(lambda x: path_filter(x.path), files))  # type: ignore
-
-        glob_filter = cast(str, path_filter)
-
-        def _filter_function(file: LazyLoadedFile) -> bool:
-            return fnmatch(file.path, glob_filter)
-
-        return list(filter(_filter_function, files))
-
     def _get_detailed_dict(self) -> types.DataRecordDetailedDict:
         return cast(types.DataRecordDetailedDict, api_client.get(f'/resources/data-records/{self.uuid}/').json())

pybiolib-1.2.174/biolib/_internal/types/file_node.py

@@ -0,0 +1,17 @@
+from .typing import TypedDict
+
+
+class FileZipMetadataDict(TypedDict):
+    header_start: int
+    size_on_disk: int
+
+
+class FileNodeDict(TypedDict):
+    dir_path: str
+    is_dir: bool
+    name: str
+    size: int
+
+
+class ZipFileNodeDict(FileNodeDict):
+    zip_meta: FileZipMetadataDict

{pybiolib-1.2.166 → pybiolib-1.2.174}/biolib/compute_node/remote_host_proxy.py

@@ -309,6 +309,15 @@ http {{
             proxy_ssl_server_name    on;
         }}
 
+        location /api/proxy/files/data-record-versions/ {{
+            proxy_pass               https://$upstream_hostname$request_uri;
+            proxy_set_header         authorization "";
+            proxy_set_header         compute-node-auth-token "{compute_node_auth_token}";
+            proxy_set_header         job-uuid "{self._job_uuid}";
+            proxy_set_header         cookie "";
+            proxy_ssl_server_name    on;
+        }}
+
         location /api/ {{
             proxy_pass               https://$upstream_hostname$request_uri;
             proxy_set_header         authorization "";

{pybiolib-1.2.166 → pybiolib-1.2.174}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "pybiolib"
-version = "1.2.166"
+version = "1.2.174"
 description = "BioLib Python Client"
 readme = "PYPI_README.md"
 license = "MIT"