enhance the functionality of _update_env

Netflix · Aug 19, 2024 · f5ed5c2 · f5ed5c2
1 parent b44f325
commit f5ed5c2
Showing 1 changed file with 7 additions and 2 deletions.
diff --git 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