diff --git a/nucleus/deploy/model_bundle.py b/nucleus/deploy/model_bundle.py index c8391f04..f5c095ac 100644 --- a/nucleus/deploy/model_bundle.py +++ b/nucleus/deploy/model_bundle.py @@ -1,11 +1,30 @@ +import inspect +from dataclasses import dataclass +from typing import Any + + +@dataclass class ModelBundle: """ Represents a ModelBundle. TODO fill this out with more than just a name potentially. """ - def __init__(self, name): - self.name = name + name: str + func_or_class: Any = None + code: str = "" def __str__(self): return f"ModelBundle(name={self.name})" + + +def create(name=None): + def decorator(func_or_class): + func_or_class.bundle_object = ModelBundle( + name=name or func_or_class.__name__, + func_or_class=func_or_class, + code=inspect.getsource(func_or_class), + ) + return func_or_class + + return decorator