pybiolib 1.2.166__py3-none-any.whl → 1.2.174__py3-none-any.whl

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())

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

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.dist-info → pybiolib-1.2.174.dist-info}/METADATA

@@ -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.dist-info → pybiolib-1.2.174.dist-info}/RECORD

@@ -1,7 +1,7 @@
 LICENSE,sha256=F2h7gf8i0agDIeWoBPXDMYScvQOz02pAWkKhTGOHaaw,1067
 PYPI_README.md,sha256=_IH7pxFiqy2bIAmaVeA-iVTyUwWRjMIlfgtUbYTtmls,368
 biolib/__init__.py,sha256=4Rfa0AJKztqkCG5D67kFgSwdOUiBTa5HkAzIOqHNREU,4431
-biolib/_data_record/data_record.py,sha256=zVAhFU1RLI1-ptoQ_l639RNwrMANXV9j75yXHvB7dtA,10950
+biolib/_data_record/data_record.py,sha256=zKvnh5T-dIVY46-kgVzMBoZ666ZhcTCFQnWvZT0D6RM,12026
 biolib/_internal/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 biolib/_internal/data_record/__init__.py,sha256=fGdME6JGRU_2VxpJbYpGXYndjN-feUkmKY4fuMyq3cg,76
 biolib/_internal/data_record/data_record.py,sha256=g_-jdy5-Zem3dthwxJj2OuQqkDGTyc-iGqN1rtYYD1A,4418
@@ -21,6 +21,7 @@ biolib/_internal/types/__init__.py,sha256=xLgOQJFh3GRtiqIJq7MaqHReZx4pp34_zcaFQ_
 biolib/_internal/types/app.py,sha256=Mz2QGD_jESX-K9JYnLWPo4YA__Q_1FQQTk9pvidCohU,118
 biolib/_internal/types/data_record.py,sha256=9r_vdhVs60YTnzU4XQFXfDrfS2P2MqD3BH2xa7lk6ck,852
 biolib/_internal/types/experiment.py,sha256=D94iBdn2nS92lRW-TOs1a2WKXJD5ZtmzL4ypggKX2ys,176
+biolib/_internal/types/file_node.py,sha256=T6BIqo662f3nwMBRqtBHYsg6YuuUaKpiokHcVjv9_ME,283
 biolib/_internal/types/resource.py,sha256=G-vPkZoe4Um6FPxsQZtRzAlbSW5sDW4NFkbjn21I3V4,372
 biolib/_internal/types/resource_version.py,sha256=sLxViYXloDDUhTDFgjegiQCj097OM1Ih1-uqlC_4ULA,174
 biolib/_internal/types/typing.py,sha256=D4EKKEe7kDx0K6lJi-H_XLtk-8w6nu2fdqn9bvzI-Xo,288
@@ -89,7 +90,7 @@ biolib/compute_node/job_worker/large_file_system.py,sha256=XXqRlVtYhs-Ji9zQGIk5K
 biolib/compute_node/job_worker/mappings.py,sha256=Z48Kg4nbcOvsT2-9o3RRikBkqflgO4XeaWxTGz-CNvI,2499
 biolib/compute_node/job_worker/utilization_reporter_thread.py,sha256=7tm5Yk9coqJ9VbEdnO86tSXI0iM0omwIyKENxdxiVXk,8575
 biolib/compute_node/job_worker/utils.py,sha256=wgxcIA8yAhUPdCwyvuuJ0JmreyWmmUoBO33vWtG60xg,1282
-biolib/compute_node/remote_host_proxy.py,sha256=aRo27m6KtNVLJm1nJSNv4ZKqdRj3E4hgaJ1m_912D28,16085
+biolib/compute_node/remote_host_proxy.py,sha256=1EZnB0WWG359Yy22xBoqRqGEmmlA12rpiG5u8B2381M,16533
 biolib/compute_node/socker_listener_thread.py,sha256=T5_UikA3MB9bD5W_dckYLPTgixh72vKUlgbBvj9dbM0,1601
 biolib/compute_node/socket_sender_thread.py,sha256=YgamPHeUm2GjMFGx8qk-99WlZhEs-kAb3q_2O6qByig,971
 biolib/compute_node/utils.py,sha256=M7i_WTyxbFM3Lri9RWZ_8FeQNYrQIWpKGLfp2I55oeY,4677
@@ -119,8 +120,8 @@ biolib/utils/cache_state.py,sha256=u256F37QSRIVwqKlbnCyzAX4EMI-kl6Dwu6qwj-Qmag,3
 biolib/utils/multipart_uploader.py,sha256=XvGP1I8tQuKhAH-QugPRoEsCi9qvbRk-DVBs5PNwwJo,8452
 biolib/utils/seq_util.py,sha256=Ozk0blGtPur_D9MwShD02r_mphyQmgZkx-lOHOwnlIM,6730
 biolib/utils/zip/remote_zip.py,sha256=0wErYlxir5921agfFeV1xVjf29l9VNgGQvNlWOlj2Yc,23232
-pybiolib-1.2.166.dist-info/LICENSE,sha256=F2h7gf8i0agDIeWoBPXDMYScvQOz02pAWkKhTGOHaaw,1067
-pybiolib-1.2.166.dist-info/METADATA,sha256=8aAKdt7iWjDaX-qQrEUsnzJuuyjE_XkMCPMvByNMYoA,1507
-pybiolib-1.2.166.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-pybiolib-1.2.166.dist-info/entry_points.txt,sha256=p6DyaP_2kctxegTX23WBznnrDi4mz6gx04O5uKtRDXg,42
-pybiolib-1.2.166.dist-info/RECORD,,
+pybiolib-1.2.174.dist-info/LICENSE,sha256=F2h7gf8i0agDIeWoBPXDMYScvQOz02pAWkKhTGOHaaw,1067
+pybiolib-1.2.174.dist-info/METADATA,sha256=whtDRAsVaAK8ZdCekzmffyOrdtKgkagocbSclm4sjrc,1507
+pybiolib-1.2.174.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+pybiolib-1.2.174.dist-info/entry_points.txt,sha256=p6DyaP_2kctxegTX23WBznnrDi4mz6gx04O5uKtRDXg,42
+pybiolib-1.2.174.dist-info/RECORD,,