From e91a6bbb6e5c040f635affd0b5d913ecb79fbaaf Mon Sep 17 00:00:00 2001 From: Robert Bradshaw Date: Mon, 11 Jul 2016 14:06:39 -0700 Subject: [PATCH] Pickle only used symbols from __main__ namespace. --- sdks/python/apache_beam/internal/pickler.py | 18 ++++++++++++++++++ .../apache_beam/internal/pickler_test.py | 3 +++ 2 files changed, 21 insertions(+) diff --git a/sdks/python/apache_beam/internal/pickler.py b/sdks/python/apache_beam/internal/pickler.py index 898e04b9a898..f427aa5eac46 100644 --- a/sdks/python/apache_beam/internal/pickler.py +++ b/sdks/python/apache_beam/internal/pickler.py @@ -159,6 +159,24 @@ def new_save_module_dict(pickler, obj): return old_save_module_dict(pickler, obj) dill.dill.save_module_dict = new_save_module_dict + + old_save_function = dill.dill.save_function + + @dill.dill.register(types.FunctionType) + def new_save_function(pickler, obj): + globs = obj.__globals__ if dill.dill.PY3 else obj.func_globals + if (dill.dill.is_dill(pickler) and globs == pickler._main.__dict__ + and not pickler._recurse): + try: + pickler._recurse = True + return old_save_function(pickler, obj) + finally: + pickler._recurse = False + else: + return old_save_function(pickler, obj) + dill.dill.save_function = new_save_function + + def _nest_dill_logging(): """Prefix all dill logging with its depth in the callstack. diff --git a/sdks/python/apache_beam/internal/pickler_test.py b/sdks/python/apache_beam/internal/pickler_test.py index 20de923ae335..bd59c73d464b 100644 --- a/sdks/python/apache_beam/internal/pickler_test.py +++ b/sdks/python/apache_beam/internal/pickler_test.py @@ -40,6 +40,9 @@ def test_lambda_with_globals(self): ['abc', 'def'], loads(dumps(module_test.get_lambda_with_globals()))('abc def')) + def test_lambda_with_main_globals(self): + self.assertEquals(unittest, loads(dumps(lambda : unittest))()) + def test_lambda_with_closure(self): """Tests that the closure of a function is preserved.""" self.assertEquals(