diff --git a/lolafect/lolaconfig.py b/lolafect/lolaconfig.py index a8265eb..5f45188 100644 --- a/lolafect/lolaconfig.py +++ b/lolafect/lolaconfig.py @@ -1,4 +1,4 @@ -from typing import List +from typing import List, Callable from prefect.storage import S3 import boto3 @@ -14,7 +14,15 @@ from lolafect.defaults import ( from lolafect.utils import S3FileReader -def _needs_env_data(method): +def _needs_env_data(method: Callable) -> Callable: + """ + Decorator to wrap methods from LolaConfig that depend on having the env + data file loaded in memory. + + :param method: the method that needs env data. + :return: the method, wrapped with a call to fetch the env data. + """ + def wrapper(self, *args, **kwargs): self.fetch_env_data(kwargs.get("s3_reader", None)) return method(self, *args, **kwargs) @@ -130,16 +138,23 @@ class LolaConfig: "ssh_jumphost": self.ENV_DATA["pt_ssh_jumphost"], } - def fetch_env_data(self, s3_reader=None): + def fetch_env_data(self, s3_reader=None) -> None: + """ + Read the env file from S3 with the default or a passed s3_reader and + store the contents in the object. + + :param s3_reader: an optional s3_reader to use instead of the default + one. + :return: None + """ if self.ENV_DATA is not None: return if s3_reader is None: s3_reader = self._s3_reader - env_data = s3_reader.read_json_from_s3_file( + self.ENV_DATA = s3_reader.read_json_from_s3_file( bucket=self.S3_BUCKET_NAME, key=self.ENV_FILE_PATH ) - self.ENV_DATA = env_data def build_lolaconfig(