diff --git a/polygon/rest/client.py b/polygon/rest/client.py index f3ff2ef2..ecc28f71 100644 --- a/polygon/rest/client.py +++ b/polygon/rest/client.py @@ -17,6 +17,15 @@ def __init__(self, auth_key: str): self._session = requests.Session() self._session.params["apiKey"] = self.auth_key + def __enter__(self): + return self + + def __exit__(self, *args): + self.close() + + def close(self): + self._session.close() + def _handle_response(self, response_type: str, endpoint: str, params: Dict[str, str]) -> Type[models.AnyDefinition]: resp: requests.Response = self._session.get(endpoint, params=params) if resp.status_code == 200: diff --git a/rest-example.py b/rest-example.py index 0b292a05..626f4660 100644 --- a/rest-example.py +++ b/rest-example.py @@ -3,10 +3,12 @@ def main(): key = "your api key" - client = RESTClient(key) - resp = client.stocks_equities_daily_open_close("AAPL", "2018-03-02") - print(f"On: {resp.from_} Apple opened at {resp.open} and closed at {resp.close}") + # RESTClient can be used as a context manager to facilitate closing the underlying http session + # https://requests.readthedocs.io/en/master/user/advanced/#session-objects + with RESTClient(key) as client: + resp = client.stocks_equities_daily_open_close("AAPL", "2018-03-02") + print(f"On: {resp.from_} Apple opened at {resp.open} and closed at {resp.close}") if __name__ == '__main__':