diff --git a/openff/qcsubmit/results/results.py b/openff/qcsubmit/results/results.py index cf2afc97..31ea8767 100644 --- a/openff/qcsubmit/results/results.py +++ b/openff/qcsubmit/results/results.py @@ -189,7 +189,9 @@ def _validate_record_types( ) @abc.abstractmethod - def to_records(self) -> List[Tuple[BaseRecord, Molecule]]: + def to_records( + self, client_kwargs: Optional[dict] = None + ) -> List[Tuple[BaseRecord, Molecule]]: """Returns the native QCPortal record objects for each of the records referenced in this collection along with a corresponding OpenFF molecule object. """ @@ -371,7 +373,9 @@ def from_server( spec_name, ) - def to_records(self) -> List[Tuple[SinglepointRecord, Molecule]]: + def to_records( + self, client_kwargs: Optional[dict] = None + ) -> List[Tuple[SinglepointRecord, Molecule]]: """Returns the native QCPortal record objects for each of the records referenced in this collection along with a corresponding OpenFF molecule object. @@ -379,9 +383,11 @@ def to_records(self) -> List[Tuple[SinglepointRecord, Molecule]]: """ records_and_molecules = [] + if client_kwargs is None: + client_kwargs = dict() for client_address, records in self.entries.items(): - client = PortalClient(client_address) + client = PortalClient(client_address, **client_kwargs) # TODO - batching/chunking (maybe in portal?) for record in records: @@ -522,7 +528,9 @@ def from_server( spec_name, ) - def to_records(self) -> List[Tuple[OptimizationRecord, Molecule]]: + def to_records( + self, client_kwargs: Optional[dict] = None + ) -> List[Tuple[OptimizationRecord, Molecule]]: """Returns the native QCPortal record objects for each of the records referenced in this collection along with a corresponding OpenFF molecule object. @@ -531,9 +539,11 @@ def to_records(self) -> List[Tuple[OptimizationRecord, Molecule]]: """ records_and_molecules = [] + if client_kwargs is None: + client_kwargs = dict() for client_address, results in self.entries.items(): - client = PortalClient(client_address) + client = PortalClient(client_address, **client_kwargs) rec_ids = [result.record_id for result in results] # Do one big request to save time @@ -791,7 +801,9 @@ def from_server( spec_name, ) - def to_records(self) -> List[Tuple[TorsiondriveRecord, Molecule]]: + def to_records( + self, client_kwargs: Optional[dict] = None + ) -> List[Tuple[TorsiondriveRecord, Molecule]]: """Returns the native QCPortal record objects for each of the records referenced in this collection along with a corresponding OpenFF molecule object. @@ -800,9 +812,11 @@ def to_records(self) -> List[Tuple[TorsiondriveRecord, Molecule]]: """ records_and_molecules = [] + if client_kwargs is None: + client_kwargs = dict() for client_address, records in self.entries.items(): - client = PortalClient(client_address) + client = PortalClient(client_address, **client_kwargs) # retrieve all torsiondrives at once, including their # minimum_optimizations