diff --git a/datamaxi/api.py b/datamaxi/api.py index 42f5c90..d004095 100644 --- a/datamaxi/api.py +++ b/datamaxi/api.py @@ -114,6 +114,15 @@ def __repr__(self): type(self).__name__, self.base_url, bool(self.api_key) ) + def close(self): + self.session.close() + + def __enter__(self): + return self + + def __exit__(self, *exc): + self.close() + def query(self, url_path, payload=None): return self.send_request("GET", url_path, payload=payload) diff --git a/datamaxi/resources/__init__.py b/datamaxi/resources/__init__.py index 1f74e4a..e32aa5f 100644 --- a/datamaxi/resources/__init__.py +++ b/datamaxi/resources/__init__.py @@ -31,7 +31,11 @@ class Datamaxi: - """Client to fetch unified data from DataMaxi+ API.""" + """Client to fetch unified data from DataMaxi+ API. + + Use as a context manager so the underlying ``requests.Session`` is + closed, or call :meth:`close` explicitly. + """ def __init__(self, api_key=None, **kwargs: Any): """Initialize the object. @@ -64,6 +68,15 @@ def __init__(self, api_key=None, **kwargs: Any): self.margin_borrow = MarginBorrow(api=api) self.index_price = IndexPrice(api=api) + def close(self): + self._api.close() + + def __enter__(self): + return self + + def __exit__(self, *exc): + self.close() + def __repr__(self): return "Datamaxi(base_url={!r}, has_key={})".format( self._api.base_url, bool(self._api.api_key) diff --git a/tests/test_api.py b/tests/test_api.py index db64ac4..f7a4613 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -118,3 +118,19 @@ def test_API_with_show_header(): with pytest.warns(DeprecationWarning, match="show_header"): client = API(show_header=True) assert client.show_header is True + + +def test_API_close_closes_session(): + """`close()` closes the underlying `requests.Session`.""" + client = API() + client.session.close = lambda: setattr(client.session, "closed", True) + client.close() + assert client.session.closed is True + + +def test_API_context_manager_closes_session(): + """`with API(...)` closes the session on exit.""" + with API() as client: + client.session.close = lambda: setattr(client.session, "closed", True) + assert client.__enter__() is client + assert client.session.closed is True diff --git a/tests/test_repr_and_lazy_pandas.py b/tests/test_repr_and_lazy_pandas.py index 30a1f9e..89e0ac2 100644 --- a/tests/test_repr_and_lazy_pandas.py +++ b/tests/test_repr_and_lazy_pandas.py @@ -34,6 +34,14 @@ def test_datamaxi_repr_and_no_key_leak(monkeypatch): assert "has_key=False" in repr(Datamaxi(base_url=BASE_URL)) +def test_datamaxi_context_manager_closes_session(): + """`with Datamaxi(...)` closes the shared session on exit (see #155).""" + with Datamaxi(api_key="secret", base_url=BASE_URL) as c: + c._api.session.close = lambda: setattr(c._api.session, "closed", True) + assert c.__enter__() is c + assert c._api.session.closed is True + + def test_importing_datamaxi_does_not_load_pandas(): # Isolated subprocess: other tests in this session load pandas, so a # same-process sys.modules check would be unreliable.