add remove existing graph arg to zntrack.Project (#544)

* add exception and remove_existing_graph

* change check

Author: PythonFZ

tests/integration/test_project.py

@@ -71,3 +71,19 @@ def test_WriteIO_no_name(tmp_path_2, assert_before_exp):
 
     assert exp2["WriteIO"].inputs == "Lorem Ipsum"
     assert exp2["WriteIO"].outputs == "Lorem Ipsum"
+
+
+def test_project_remove_graph(proj_path):
+    with zntrack.Project() as project:
+        node = WriteIO(inputs="Hello World")
+    project.run()
+    node.load()
+    assert node.outputs == "Hello World"
+
+    with zntrack.Project(remove_existing_graph=True) as project:
+        node2 = WriteIO(inputs="Lorem Ipsum", name="node2")
+    project.run()
+    node2.load()
+    assert node2.outputs == "Lorem Ipsum"
+    with pytest.raises(zntrack.exceptions.NodeNotAvailableError):
+        node.load()

zntrack/__init__.py

@@ -4,7 +4,7 @@
 """
 import importlib.metadata
 
-from zntrack import tools
+from zntrack import exceptions, tools
 from zntrack.core.node import Node
 from zntrack.core.nodify import NodeConfig, nodify
 from zntrack.fields import Field, FieldGroup, LazyField, dvc, meta, zn
@@ -28,4 +28,5 @@
     "nodify",
     "NodeConfig",
     "tools",
+    "exceptions",
 ]

zntrack/core/node.py

@@ -14,6 +14,7 @@
 import zninit
 import znjson
 
+from zntrack import exceptions
 from zntrack.notebooks.jupyter import jupyter_class_to_file
 from zntrack.utils import NodeStatusResults, deprecated, module_handler, run_dvc_cmd
 from zntrack.utils.config import config
@@ -162,10 +163,13 @@ def load(self, lazy: bool = None) -> None:
 
         kwargs = {} if lazy is None else {"lazy": lazy}
         self.state.loaded = True  # we assume loading will be successful.
-        with config.updated_config(**kwargs):
-            # TODO: it would be much nicer not to use a global config object here.
-            for attr in zninit.get_descriptors(Field, self=self):
-                attr.load(self)
+        try:
+            with config.updated_config(**kwargs):
+                # TODO: it would be much nicer not to use a global config object here.
+                for attr in zninit.get_descriptors(Field, self=self):
+                    attr.load(self)
+        except KeyError as err:
+            raise exceptions.NodeNotAvailableError(self) from err
 
         # TODO: documentation about _post_init and _post_load_ and when they are called
         self._post_load_()

zntrack/exceptions/__init__.py

@@ -0,0 +1,19 @@
+"""All ZnTrack exceptions."""
+
+
+class NodeNotAvailableError(Exception):
+    """Raised when a node is not available."""
+
+    def __init__(self, arg):
+        """Initialize the exception.
+
+        Parameters
+        ----------
+        arg : str|Node
+            Custom Error message or Node that is not available.
+        """
+        if isinstance(arg, str):
+            super().__init__(arg)
+        else:
+            # assume arg is a Node
+            super().__init__(f"Node {arg.name} is not available.")

zntrack/project/zntrack_project.py

@@ -96,19 +96,29 @@ def _initalize():
 class Project(_ProjectBase):
     """The ZnTrack Project class."""
 
-    def __init__(self, initialize: bool = True) -> None:
+    def __init__(
+        self, initialize: bool = True, remove_existing_graph: bool = False
+    ) -> None:
         """Initialize the Project.
 
         Attributes
         ----------
         initialize : bool, default = True
             If True, initialize a git repository and a dvc repository.
+        remove_existing_graph : bool, default = False
+            If True, remove 'dvc.yaml', 'zntrack.json' and 'params.yaml'
+              before writing new nodes.
         """
         # TODO maybe it is not a good idea to base everything on the DiGraph class.
         #  It seems to call some class methods
         super().__init__()
         if initialize:
             _initalize()
+        if remove_existing_graph:
+            # we remove the files that typically contain the graph definition
+            pathlib.Path("zntrack.json").unlink(missing_ok=True)
+            pathlib.Path("dvc.yaml").unlink(missing_ok=True)
+            pathlib.Path("params.yaml").unlink(missing_ok=True)
 
     def create_branch(self, name: str) -> "Branch":
         """Create a branch in the project."""