nicer output

This commit is contained in:
James Turk 2023-05-08 03:14:45 -05:00
parent 3366cb2611
commit 3eb14f9556

View File

@ -19,6 +19,21 @@ def get_sha512(filename: str) -> str:
return hashlib.sha512(file.read()).hexdigest() return hashlib.sha512(file.read()).hexdigest()
@dataclass(frozen=True, eq=True)
class Transform:
name: str
transform_func: callable
error_map: dict[tuple, str]
def if_cond_true(data_cond_tup: tuple[dict, bool]) -> dict | None:
return data_cond_tup[0] if data_cond_tup[1] else None
def if_cond_false(data_cond_tup: tuple[dict, bool]) -> dict | None:
return data_cond_tup[0] if not data_cond_tup[1] else None
class Recipe: class Recipe:
def __init__(self, name, db_name="beakers.db"): def __init__(self, name, db_name="beakers.db"):
self.name = name self.name = name
@ -46,10 +61,23 @@ class Recipe:
from_beaker: str, from_beaker: str,
to_beaker: str, to_beaker: str,
transform_func: callable, transform_func: callable,
*,
name=None,
error_map: dict[tuple, str] | None = None, error_map: dict[tuple, str] | None = None,
) -> None: ) -> None:
if name is None:
name = transform_func.__name__
if name == "<lambda>":
name = "λ"
transform = Transform(
name=name,
transform_func=transform_func,
error_map=error_map or {},
)
self.graph.add_edge( self.graph.add_edge(
from_beaker, to_beaker, transform=transform_func, error_map=error_map from_beaker,
to_beaker,
transform=transform,
) )
def add_conditional( def add_conditional(
@ -60,28 +88,28 @@ class Recipe:
if_false: str, if_false: str,
) -> None: ) -> None:
# first add a transform to evaluate the conditional # first add a transform to evaluate the conditional
if condition_func.__name__ == "<lambda>":
cond_name = f"cond-{from_beaker}"
else:
cond_name = f"cond-{from_beaker}-{condition_func.__name__}" cond_name = f"cond-{from_beaker}-{condition_func.__name__}"
self.add_beaker(cond_name, temp=True) self.add_beaker(cond_name, temp=True)
self.add_transform( self.add_transform(
from_beaker, from_beaker,
cond_name, cond_name,
lambda data: (data, condition_func(data)), lambda data: (data, condition_func(data)),
name=cond_name,
) )
# then add two filtered paths that remove the condition result # then add two filtered paths that remove the condition result
self.graph.add_edge( self.add_transform(
cond_name, cond_name,
if_true, if_true,
transform=lambda data_cond_tup: data_cond_tup[0] if_cond_true,
if data_cond_tup[1]
else None,
) )
self.graph.add_edge( self.add_transform(
cond_name, cond_name,
if_false, if_false,
transform=lambda data_cond_tup: data_cond_tup[0] if_cond_false,
if not data_cond_tup[1]
else None,
) )
def get_metadata(self, table_name) -> dict: def get_metadata(self, table_name) -> dict:
@ -150,7 +178,13 @@ class Recipe:
fg=typer.colors.GREEN if lb else typer.colors.YELLOW, fg=typer.colors.GREEN if lb else typer.colors.YELLOW,
) )
for from_b, to_b, edge in self.graph.out_edges(node, data=True): for from_b, to_b, edge in self.graph.out_edges(node, data=True):
print(f" {from_b} -> {to_b} ({edge})") name = edge["transform"].name
print(f" {from_b} -({name})-> {to_b}")
for k, v in edge["transform"].error_map.items():
typer.secho(
f" {' '.join(c.__name__ for c in k)} -> {v}",
fg=typer.colors.RED,
)
def run_once( def run_once(
self, start_beaker: str | None = None, end_beaker: str | None = None self, start_beaker: str | None = None, end_beaker: str | None = None
@ -178,21 +212,25 @@ class Recipe:
edges = self.graph.out_edges(node, data=True) edges = self.graph.out_edges(node, data=True)
for from_b, to_b, edge in edges: for from_b, to_b, edge in edges:
if "transform" not in edge: transform = edge["transform"]
raise ValueError(f"unknown edge: {edge}")
error_map = edge.get("error_map", {})
from_beaker = self.beakers[from_b] from_beaker = self.beakers[from_b]
to_beaker = self.beakers[to_b] to_beaker = self.beakers[to_b]
func = edge["transform"] log.info(
log.info("transform", from_b=from_b, to_b=to_b, items=len(from_beaker)) "transform",
from_b=from_b,
to_b=to_b,
items=len(from_beaker),
transform=edge["transform"].name,
)
# convert coroutine to function # convert coroutine to function
if inspect.iscoroutinefunction(func): if inspect.iscoroutinefunction(transform.transform_func):
t_func = lambda x: loop.run_until_complete(func(x)) t_func = lambda x: loop.run_until_complete(
transform.transform_func(x)
)
else: else:
t_func = func t_func = transform.transform_func
for id, item in from_beaker.items(): for id, item in from_beaker.items():
try: try:
@ -200,7 +238,10 @@ class Recipe:
if transformed: if transformed:
to_beaker.add_item(transformed, id) to_beaker.add_item(transformed, id)
except Exception as e: except Exception as e:
for error_types, error_beaker_name in error_map.items(): for (
error_types,
error_beaker_name,
) in transform.error_map.items():
if isinstance(e, error_types): if isinstance(e, error_types):
error_beaker = self.beakers[error_beaker_name] error_beaker = self.beakers[error_beaker_name]
error_beaker.add_item( error_beaker.add_item(