diff --git a/src/bmkg/api/api.py b/src/bmkg/api/api.py index eb37f9c..e126065 100644 --- a/src/bmkg/api/api.py +++ b/src/bmkg/api/api.py @@ -1,3 +1,7 @@ +from traceback import TracebackException +from types import TracebackType +from typing import Self + from aiohttp import ClientSession __all__ = ["API"] @@ -12,3 +16,17 @@ def __init__(self, session: ClientSession | None = None) -> None: self._session = ( session if session is not None else ClientSession("https://data.bmkg.go.id") ) + + async def __aenter__(self) -> Self: + return self + + async def __aexit__( + self, + exc_type: Exception, + exc_val: TracebackException, + traceback: TracebackType, + ) -> None: + await self.close() + + async def close(self) -> None: + await self._session.close() diff --git a/src/bmkg/api/bmkg.py b/src/bmkg/api/bmkg.py index a6afc21..392d959 100644 --- a/src/bmkg/api/bmkg.py +++ b/src/bmkg/api/bmkg.py @@ -1,6 +1,4 @@ -from traceback import TracebackException -from types import TracebackType -from typing import Self +from aiohttp import ClientSession from .api import API from .earthquake import Earthquake @@ -18,22 +16,8 @@ class BMKG(API): weather_forecast (WeatherForecast): weather forecast api interface. """ - def __init__(self) -> None: - API.__init__(self) + def __init__(self, session: ClientSession | None = None) -> None: + API.__init__(self, session) self.earthquake = Earthquake(self._session) self.weather_forecast = WeatherForecast(self._session) - - async def __aenter__(self) -> Self: - return self - - async def __aexit__( - self, - exc_type: Exception, - exc_val: TracebackException, - traceback: TracebackType, - ) -> None: - await self.close() - - async def close(self) -> None: - await self._session.close() diff --git a/src/bmkg/api/earthquake.py b/src/bmkg/api/earthquake.py index 9ef789a..8a02776 100644 --- a/src/bmkg/api/earthquake.py +++ b/src/bmkg/api/earthquake.py @@ -43,7 +43,7 @@ async def get_latest_earthquake(self) -> schemas.LatestEarthquake: latest_earthquake = parse_latest_earthquake_data(await response.json()) # type: ignore latest_earthquake.shakemap = Shakemap( - self._session, latest_earthquake.shakemap.file_name + latest_earthquake.shakemap.file_name, self._session ) return latest_earthquake diff --git a/src/bmkg/api/shakemap.py b/src/bmkg/api/shakemap.py index 0cb2c90..177d25c 100644 --- a/src/bmkg/api/shakemap.py +++ b/src/bmkg/api/shakemap.py @@ -13,7 +13,7 @@ class Shakemap(API, schemas.Shakemap): url = "/DataMKG/TEWS" - def __init__(self, session: ClientSession, file_name: str) -> None: + def __init__(self, file_name: str, session: ClientSession | None = None) -> None: API.__init__(self, session) schemas.Shakemap.__init__(self, file_name)