fix: sync API/Datamaxi close() + context manager, fixes #155

mirrors AsyncAPI/AsyncDatamaxi aclose/__aenter__/__aexit__

Author: martinkersner

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)
 

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)

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

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.