diff --git a/tests/integration/test_single_node.py b/tests/integration/test_single_node.py index f15c4dbd..31e2c020 100644 --- a/tests/integration/test_single_node.py +++ b/tests/integration/test_single_node.py @@ -15,18 +15,25 @@ def run(self): self.c = self.a + self.b +class AddNumbersMetrics(AddNumbers): + c = zntrack.zn.metrics() + + @pytest.mark.parametrize("eager", [True, False]) def test_AddNumbers(proj_path, eager): with zntrack.Project() as project: - add_numbers = AddNumbers(a=1, b=2) + add_numbers = AddNumbersMetrics(a=1, b=2) assert not add_numbers.state.loaded project.run(eager=eager) if not eager: add_numbers.load() + assert add_numbers.state.params == {"a": 1, "b": 2} + assert add_numbers.c == 3 assert add_numbers.state.loaded + assert add_numbers.name == "AddNumbersMetrics" def test_AddNumbers_remove_params(proj_path): diff --git a/zntrack/core/node.py b/zntrack/core/node.py index 8a503fc2..b186459a 100644 --- a/zntrack/core/node.py +++ b/zntrack/core/node.py @@ -55,6 +55,8 @@ class NodeStatus: remote: str = None rev: str = None + parent: Node = dataclasses.field(default=None, repr=False) + def get_file_system(self) -> dvc.api.DVCFileSystem: """Get the file system of the Node.""" log.warning("Deprecated. Use 'state.fs' instead.") @@ -74,6 +76,21 @@ def fs(self) -> dvc.api.DVCFileSystem: time.sleep(0.1) raise dvc.utils.strictyaml.YAMLValidationError + @property + def params(self) -> dict: + """Get the parameters of the Node.""" + params = dvc.api.params_show( + stages=self.parent.name, + repo=self.remote, + rev=self.rev, + ) + return params.get(self.parent.name, {}) + + @property + def metrics(self) -> dict: + """Get the metrics of the Node.""" + raise ValueError("Not implemented yet.") + class _NameDescriptor(zninit.Descriptor): """A descriptor for the name attribute.""" @@ -155,7 +172,7 @@ def _init_descriptors_(self): def state(self) -> NodeStatus: """Get the state of the node.""" if self._state is None: - self._state = NodeStatus(False, NodeStatusResults.UNKNOWN) + self._state = NodeStatus(False, NodeStatusResults.UNKNOWN, parent=self) return self._state @property