counterweight/lolafect
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Some docstrings, typing and formatting.

Browse source
This commit is contained in:
Pablo Martin 2023-01-16 13:33:45 +01:00
parent 4f22ad0172
commit 610288aafe
1 changed files with 20 additions and 5 deletions
Download patch file Download diff file Expand all files Collapse all files

25
lolafect/lolaconfig.py
View file

@ -1,4 +1,4 @@
from typing import List from typing import List, Callable
from prefect.storage import S3 from prefect.storage import S3
import boto3 import boto3
@ -14,7 +14,15 @@ from lolafect.defaults import (
from lolafect.utils import S3FileReader 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): def wrapper(self, *args, **kwargs):
self.fetch_env_data(kwargs.get("s3_reader", None)) self.fetch_env_data(kwargs.get("s3_reader", None))
return method(self, *args, **kwargs) return method(self, *args, **kwargs)
@ -130,16 +138,23 @@ class LolaConfig:
"ssh_jumphost": self.ENV_DATA["pt_ssh_jumphost"], "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: if self.ENV_DATA is not None:
return return
if s3_reader is None: if s3_reader is None:
s3_reader = self._s3_reader 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 bucket=self.S3_BUCKET_NAME, key=self.ENV_FILE_PATH
) )
self.ENV_DATA = env_data
def build_lolaconfig( def build_lolaconfig(