diff --git a/heroku_audit/client.py b/heroku_audit/client.py index 115bf8f..0a55252 100644 --- a/heroku_audit/client.py +++ b/heroku_audit/client.py @@ -1,12 +1,35 @@ import os import sys +from typing import Any, cast import heroku3 +from heroku3.core import Heroku from rich import print from rich.text import Text -try: - heroku = heroku3.from_key(os.environ["HEROKU_API_KEY"]) -except KeyError: - print(Text("Please set $HEROKU_API_KEY to a valid Heroku API key.", style="red")) - sys.exit(1) +__all__ = ["heroku"] + + +class LazyHerokuWrapper: + """ + A lazy heroku wrapper which only requires an API key when it's used + """ + + _heroku = None + + def __get__(self, obj: Any, objtype: Any = None) -> Heroku: + if self._heroku is None: + try: + self._heroku = heroku3.from_key(os.environ["HEROKU_API_KEY"]) + except KeyError: + print( + Text( + "Please set $HEROKU_API_KEY to a valid Heroku API key.", + style="red", + ) + ) + sys.exit(1) + return self._heroku + + +heroku = cast(Heroku, LazyHerokuWrapper()) diff --git a/stubs/heroku3/__init__.pyi b/stubs/heroku3/__init__.pyi index 09a070e..cb27262 100644 --- a/stubs/heroku3/__init__.pyi +++ b/stubs/heroku3/__init__.pyi @@ -1,11 +1,3 @@ -from requests import Session - -from .models.app import App - -class Heroku: - _session: Session - - def apps(self) -> list[App]: ... - def app(self, id_or_name: str) -> App: ... +from .core import Heroku def from_key(api_key: str) -> Heroku: ... diff --git a/stubs/heroku3/core.pyi b/stubs/heroku3/core.pyi new file mode 100644 index 0000000..064c4b0 --- /dev/null +++ b/stubs/heroku3/core.pyi @@ -0,0 +1,9 @@ +from requests import Session + +from .models.app import App + +class Heroku: + _session: Session + + def apps(self) -> list[App]: ... + def app(self, id_or_name: str) -> App: ...