From cda9064406527466f8b4575f9ed5b7a28eda80f8 Mon Sep 17 00:00:00 2001
From: Ivan Khomyakov <ivan.khomyakov@scale.com>
Date: Tue, 22 Feb 2022 12:53:27 -0800
Subject: [PATCH 1/3] add decorator in model bundle

---
 nucleus/deploy/model_bundle.py | 23 ++++++++++++++++++++---
 1 file changed, 20 insertions(+), 3 deletions(-)

diff --git a/nucleus/deploy/model_bundle.py b/nucleus/deploy/model_bundle.py
index c8391f04..1f37d0d6 100644
--- a/nucleus/deploy/model_bundle.py
+++ b/nucleus/deploy/model_bundle.py
@@ -1,11 +1,28 @@
+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

From 1a3b0570302565fb71db861fb77c82d00eba082a Mon Sep 17 00:00:00 2001
From: Ivan Khomyakov <ivan.khomyakov@scale.com>
Date: Tue, 22 Feb 2022 14:39:49 -0800
Subject: [PATCH 2/3] black

---
 nucleus/deploy/model_bundle.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/nucleus/deploy/model_bundle.py b/nucleus/deploy/model_bundle.py
index 1f37d0d6..6534b3f8 100644
--- a/nucleus/deploy/model_bundle.py
+++ b/nucleus/deploy/model_bundle.py
@@ -9,6 +9,7 @@ class ModelBundle:
     Represents a ModelBundle.
     TODO fill this out with more than just a name potentially.
     """
+
     name: str
     func_or_class: Any = None
     code: str = ""
@@ -25,4 +26,5 @@ def decorator(func_or_class):
             code=inspect.getsource(func_or_class),
         )
         return func_or_class
+
     return decorator

From a4fe3545f2a95dc59a9d2ea14a4fa2877dc52d6f Mon Sep 17 00:00:00 2001
From: Ivan Khomyakov <ivan.khomyakov@scale.com>
Date: Tue, 22 Feb 2022 15:13:07 -0800
Subject: [PATCH 3/3] linter

---
 nucleus/deploy/model_bundle.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/nucleus/deploy/model_bundle.py b/nucleus/deploy/model_bundle.py
index 6534b3f8..f5c095ac 100644
--- a/nucleus/deploy/model_bundle.py
+++ b/nucleus/deploy/model_bundle.py
@@ -20,7 +20,7 @@ def __str__(self):
 
 def create(name=None):
     def decorator(func_or_class):
-        func_or_class._bundle_object = ModelBundle(
+        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),