From f5ed5c2d603912731320bacbc25038b829cfaea2 Mon Sep 17 00:00:00 2001 From: Madhur Tandon Date: Mon, 19 Aug 2024 15:43:20 +0530 Subject: [PATCH] enhance the functionality of _update_env --- metaflow/metaflow_current.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/metaflow/metaflow_current.py b/metaflow/metaflow_current.py index ecd9730ebb1..a6690cb2a6e 100644 --- a/metaflow/metaflow_current.py +++ b/metaflow/metaflow_current.py @@ -1,5 +1,5 @@ from collections import namedtuple -import os +import functools from typing import Any, Optional, TYPE_CHECKING from metaflow.metaflow_config import TEMPDIR @@ -62,7 +62,12 @@ def _set_env( def _update_env(self, env): for k, v in env.items(): - setattr(self.__class__, k, property(fget=lambda _, v=v: v)) + if isinstance(v, property): + setattr(self.__class__, k, v) + elif callable(v): + setattr(self, k, functools.partial(v, self)) + else: + setattr(self.__class__, k, property(fget=lambda _, v=v: v)) def __contains__(self, key: str): return getattr(self, key, None) is not None