Skip to content

Commit

Permalink
feat: add context manager support for bentoml.client (#3402)
Browse files Browse the repository at this point in the history
Address #3288
  • Loading branch information
y1450 committed Jan 13, 2023
1 parent 75daa5a commit 76cb310
Showing 1 changed file with 25 additions and 0 deletions.
25 changes: 25 additions & 0 deletions src/bentoml/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,9 @@
from .exceptions import BentoMLException
from ._internal.service.inference_api import InferenceAPI

if t.TYPE_CHECKING:
from types import TracebackType


class Client(ABC):
server_url: str
Expand Down Expand Up @@ -85,6 +88,28 @@ def wait_until_server_is_ready(host: str, port: int, timeout: int) -> None:
raise TimeoutError("The server took too long to get ready")
time.sleep(1)

def __enter__(self):
return self

def __exit__(
self,
exc_type: type[BaseException] | None,
exc_value: BaseException | None,
traceback: TracebackType | None,
) -> bool | None:
pass

async def __aenter__(self):
return self

async def __aexit__(
self,
exc_type: type[BaseException] | None,
exc_value: BaseException | None,
traceback: TracebackType | None,
) -> bool | None:
pass

@staticmethod
def from_url(server_url: str) -> Client:
server_url = server_url if "://" in server_url else "http://" + server_url
Expand Down

0 comments on commit 76cb310

Please sign in to comment.