Add support for class attributes and bound methods to escape hatch (#2775)

Update the escape hatch to allow escaping the following components:
- Class-level attributes
- Also fixed a bug where falsy enum values (e.g. 0 and “”) would throw a
RuntimeError since they were considered None
- Bound methods

Author: malder08

metaflow/plugins/env_escape/client.py

@@ -378,16 +378,16 @@ def name_to_parent_name(name):
         else:
             lookup_name = name
 
-        if name == "function":
-            # Special handling of pickled functions. We create a new class that
+        if name in ("function", "method"):
+            # Special handling of pickled functions and methods. We create a new class that
             # simply has a __call__ method that will forward things back to
             # the server side.
             if obj_id is None:
-                raise RuntimeError("Local function unpickling without an object ID")
+                raise RuntimeError("Local %s unpickling without an object ID" % name)
             if obj_id not in self._proxied_standalone_functions:
                 self._proxied_standalone_functions[obj_id] = create_class(
                     self,
-                    "__main__.__function_%s" % obj_id,
+                    "__main__.__%s_%s" % (name, obj_id),
                     {},
                     {},
                     {},

metaflow/plugins/env_escape/configurations/emulate_test_lib/server_mappings.py

@@ -17,6 +17,8 @@
         "ChildClass": lib.ChildClass,
         "ExceptionAndClass": lib.ExceptionAndClass,
         "ExceptionAndClassChild": lib.ExceptionAndClassChild,
+        "TestIntEnum": lib.TestIntEnum,
+        "TestStrEnum": lib.TestStrEnum,
     }
 }
 

metaflow/plugins/env_escape/configurations/test_lib_impl/test_lib.py

@@ -1,7 +1,20 @@
 import functools
+from enum import Enum, IntEnum
 from html.parser import HTMLParser
 
 
+class TestIntEnum(IntEnum):
+    ZERO = 0
+    ONE = 1
+    TWO = 2
+
+
+class TestStrEnum(str, Enum):
+    EMPTY = ""
+    FOO = "foo"
+    BAR = "bar"
+
+
 class MyBaseException(Exception):
     pass
 
@@ -106,6 +119,9 @@ def weird_indirection(self, name):
     def returnChild(self):
         return ChildClass()
 
+    def get_bound_method(self):
+        return self.print_value
+
     def raiseOrReturnValueError(self, doRaise=False):
         if doRaise:
             raise ValueError("I raised")

metaflow/plugins/env_escape/consts.py

@@ -40,6 +40,7 @@
 OP_INIT = 17
 OP_CALLONCLASS = 18
 OP_SUBCLASSCHECK = 19
+OP_GETCLASSATTR = 20
 
 # Control messages
 CONTROL_SHUTDOWN = 1

metaflow/plugins/env_escape/server.py

@@ -37,6 +37,7 @@
     OP_SETVAL,
     OP_INIT,
     OP_SUBCLASSCHECK,
+    OP_GETCLASSATTR,
     VALUE_LOCAL,
     VALUE_REMOTE,
     CONTROL_GETEXPORTS,
@@ -96,11 +97,19 @@ def __init__(self, config_dir, max_pickle_version):
             {v: k for k, v in self._proxied_types.items()}
         )
 
-        # We will also proxy functions from objects as needed. This is useful
+        # We will also proxy functions and methods from objects as needed. This is useful
         # for defaultdict for example since the `default_factory` function is a
-        # lambda that needs to be transferred.
+        # lambda that needs to be transferred. Methods are also proxied to support
+        # bound methods returned from functions (e.g., evaluator patterns).
         self._class_types_to_names[type(lambda x: x)] = "function"
 
+        # Register method type for bound methods
+        class _TempClass:
+            def _temp_method(self):
+                pass
+
+        self._class_types_to_names[type(_TempClass()._temp_method)] = "method"
+
         # Update all alias information
         for base_name, aliases in itertools.chain(
             a1.items(), a2.items(), a3.items(), a4.items()
@@ -257,6 +266,7 @@ def __init__(self, config_dir, max_pickle_version):
             OP_SETVAL: self._handle_setval,
             OP_INIT: self._handle_init,
             OP_SUBCLASSCHECK: self._handle_subclasscheck,
+            OP_GETCLASSATTR: self._handle_getclassattr,
         }
 
         self._local_objects = {}
@@ -391,9 +401,9 @@ def pickle_object(self, obj):
     def unpickle_object(self, obj):
         if (not isinstance(obj, ObjReference)) or obj.value_type != VALUE_LOCAL:
             raise ValueError("Invalid transferred object: %s" % str(obj))
-        obj = self._local_objects.get(obj.identifier)
-        if obj:
-            return obj
+        result = self._local_objects.get(obj.identifier)
+        if result is not None:
+            return result
         raise ValueError("Invalid object -- id %s not known" % obj.identifier)
 
     @staticmethod
@@ -527,6 +537,15 @@ def _handle_subclasscheck(self, target, class_name, otherclass_name, reverse=Fal
             return issubclass(class_type, getattr(sys.modules[sub_module], sub_name))
         return issubclass(getattr(sys.modules[sub_module], sub_name), class_type)
 
+    def _handle_getclassattr(self, target, class_name, attr_name):
+        # Handle class-level attribute access like EnumClass.MEMBER
+        class_type = self._known_classes.get(class_name)
+        if class_type is None:
+            class_type = self._proxied_types.get(class_name)
+        if class_type is None:
+            raise ValueError("Unknown class %s" % class_name)
+        return getattr(class_type, attr_name)
+
 
 if __name__ == "__main__":
     max_pickle_version = int(sys.argv[1])

metaflow/plugins/env_escape/stub.py

@@ -17,6 +17,7 @@
     OP_DIR,
     OP_INIT,
     OP_SUBCLASSCHECK,
+    OP_GETCLASSATTR,
 )
 
 from .exception_transferer import ExceptionMetaClass
@@ -311,6 +312,14 @@ def __instancecheck__(cls, instance):
         # Goes to __subclasscheck__ above
         return cls.__subclasscheck__(type(instance))
 
+    def __getattr__(cls, name):
+        # This handles class-level attribute access like EnumClass.MEMBER
+        # When accessing an attribute on the stub class itself (not an instance),
+        # forward the request to the server to get the class attribute.
+        return cls.___class_connection___.stub_request(
+            None, OP_GETCLASSATTR, cls.___class_remote_class_name___, name
+        )
+
 
 class MetaExceptionWithConnection(StubMetaClass, ExceptionMetaClass):
     def __new__(cls, class_name, base_classes, class_dict, connection):

test/env_escape/example.py

@@ -137,6 +137,27 @@ def run_test(through_escape=False):
     # assert issubclass(test.ChildClass, HTMLParser)
     assert issubclass(test.ChildClass, object)
 
+    print("-- Test enum class attribute access --")
+    assert test.TestIntEnum.ZERO == 0  # Also tests falsy value (0)
+    assert test.TestIntEnum.ONE == 1
+    assert test.TestIntEnum.TWO == 2
+    assert test.TestStrEnum.EMPTY == ""  # Also tests falsy value (empty string)
+    assert test.TestStrEnum.FOO == "foo"
+    assert test.TestStrEnum.BAR == "bar"
+
+    print("-- Test bound method returns --")
+    o1_for_method = test.TestClass1(5)
+
+    # Get a bound method (not functools.partial)
+    bound_method = o1_for_method.get_bound_method()
+
+    # Verify it's callable
+    assert callable(bound_method)
+
+    # Call the bound method - returns the raw value (no override applied)
+    result = bound_method()
+    assert result == 5, f"Expected 5, got {result}"
+
     print("-- Test exceptions --")
     # Non proxied exceptions can't be returned as objects
     try: