Bump version to 0.4.2

- traps pydantic import just in case of versioning error
- formatting
- verbose options
- predict fixes
- stricter typing fixes

Author: markgeejw

openprotein/api/align.py

@@ -1,5 +1,5 @@
 from typing import Iterator, Optional, List, BinaryIO, Literal, Union
-from pydantic import BaseModel, Field, validator, root_validator
+from openprotein.pydantic import BaseModel, Field, validator, root_validator
 from enum import Enum
 from io import BytesIO
 import random

openprotein/api/data.py

@@ -1,6 +1,6 @@
 import pandas as pd
-import pydantic
-from pydantic import BaseModel
+import openprotein.pydantic as pydantic
+from openprotein.pydantic import BaseModel
 from typing import Optional, List, Union
 from datetime import datetime
 from io import BytesIO
@@ -468,7 +468,7 @@ def create(
         metadata.sequence_length = len(table["sequence"].values[0])
         return AssayDataset(self.session, metadata)
 
-    def get(self, assay_id: str) -> AssayMetadata:
+    def get(self, assay_id: str, verbose: bool = False) -> AssayMetadata:
         """
         Get an assay dataset by its ID.
 

openprotein/api/design.py

@@ -9,7 +9,7 @@
 
 from openprotein.errors import APIError
 from openprotein.futures import FutureFactory, FutureBase
-from pydantic import BaseModel, Field, validator
+from openprotein.pydantic import BaseModel, Field, validator
 from datetime import datetime
 import re
 

openprotein/api/embedding.py

@@ -11,7 +11,7 @@
 )
 from openprotein.futures import FutureBase, FutureFactory
 
-from pydantic import BaseModel, parse_obj_as
+from openprotein.pydantic import BaseModel, parse_obj_as
 import numpy as np
 from typing import Optional, List, Union, Any
 import io
@@ -247,7 +247,7 @@ def __init__(
 
     def get(self, verbose=False) -> List:
         return super().get(verbose=verbose)
-    
+
     @property
     def sequences(self):
         if self._sequences is None:
@@ -305,9 +305,7 @@ def embedding_model_post(
     """
     endpoint = PATH_PREFIX + f"/models/{model_id}/embed"
 
-    sequences_unicode = [
-        (s if isinstance(s, str) else s.decode()) for s in sequences 
-    ]
+    sequences_unicode = [(s if isinstance(s, str) else s.decode()) for s in sequences]
     body = {
         "sequences": sequences_unicode,
     }
@@ -345,9 +343,7 @@ def embedding_model_logits_post(
     """
     endpoint = PATH_PREFIX + f"/models/{model_id}/logits"
 
-    sequences_unicode = [
-        (s if isinstance(s, str) else s.decode()) for s in sequences 
-    ]
+    sequences_unicode = [(s if isinstance(s, str) else s.decode()) for s in sequences]
     body = {
         "sequences": sequences_unicode,
     }
@@ -385,9 +381,7 @@ def embedding_model_attn_post(
     """
     endpoint = PATH_PREFIX + f"/models/{model_id}/attn"
 
-    sequences_unicode = [
-        (s if isinstance(s, str) else s.decode()) for s in sequences 
-    ]
+    sequences_unicode = [(s if isinstance(s, str) else s.decode()) for s in sequences]
     body = {
         "sequences": sequences_unicode,
     }
@@ -500,9 +494,7 @@ def svd_embed_post(session: APISession, svd_id: str, sequences: List[bytes]) ->
     """
     endpoint = PATH_PREFIX + f"/svd/{svd_id}/embed"
 
-    sequences_unicode = [
-        (s if isinstance(s, str) else s.decode()) for s in sequences 
-    ]
+    sequences_unicode = [(s if isinstance(s, str) else s.decode()) for s in sequences]
     body = {
         "sequences": sequences_unicode,
     }
@@ -715,7 +707,7 @@ def get_job(self) -> Job:
         """Get job associated with this SVD model"""
         return job_get(self.session, self.id)
 
-    def get(self):
+    def get(self, verbose: bool = False):
         # overload for AsyncJobFuture
         return self
 
@@ -963,7 +955,7 @@ def fit_svd(
         sequences: List[bytes],
         n_components: int = 1024,
         reduction: Optional[str] = None,
-    ) -> SVDModel:
+    ) -> SVDModel:  # type: ignore
         """
         Fit an SVD on the embedding results of this model. 
 

openprotein/api/fold.py

@@ -3,7 +3,7 @@
 import openprotein.config as config
 from openprotein.api.embedding import ModelMetadata
 from openprotein.api.align import validate_msa, MSAFuture
-import pydantic
+import openprotein.pydantic as pydantic
 from typing import Optional, List, Union, Tuple
 from openprotein.futures import FutureBase, FutureFactory
 from abc import ABC, abstractmethod

openprotein/api/jobs.py

@@ -1,12 +1,12 @@
 # Jobs and job centric flows
 
 
-from typing import List, Union
+from typing import List, Union, Optional
 import concurrent.futures
 import time
 
 import tqdm
-import pydantic
+import openprotein.pydantic as pydantic
 
 from openprotein.base import APISession
 import openprotein.config as config
@@ -105,7 +105,7 @@ def list(
             more_recent_than=more_recent_than,
         )
 
-    def get(self, job_id) -> Job:
+    def get(self, job_id: str, verbose: bool = False) -> Job:
         """get Job by ID"""
         return load_job(self.session, job_id)
         # return job_get(self.session, job_id)
@@ -150,7 +150,7 @@ def done(self):
     def cancelled(self):
         return self.job.cancelled()
 
-    def get(self, verbose=False):
+    def get(self, verbose: bool = False):
         raise NotImplementedError()
 
     def wait_until_done(
@@ -176,7 +176,7 @@ def wait_until_done(
     def wait(
         self,
         interval: int = config.POLLING_INTERVAL,
-        timeout: int = None,
+        timeout: Optional[int] = None,
         verbose: bool = False,
     ):
         """
@@ -195,7 +195,7 @@ def wait(
             self.session, interval=interval, timeout=timeout, verbose=verbose
         )
         self.job = job
-        return self.get(verbose=verbose)
+        return self.get()
 
 
 class StreamingAsyncJobFuture(AsyncJobFuture):

openprotein/api/poet.py

@@ -1,5 +1,5 @@
 from typing import Iterator, Optional, List, Literal, Dict
-from pydantic import BaseModel, validator
+from openprotein.pydantic import BaseModel, validator
 from io import BytesIO
 import random
 import requests