From 4244d3b2e22a4aa3abff81f80cf45fa6e40934ef Mon Sep 17 00:00:00 2001
From: Leonardo Gama <leonardo.reis.gama@gmail.com>
Date: Wed, 14 Sep 2022 09:53:23 -0300
Subject: [PATCH] fix _stash_modules() arguments

---
 dill/_dill.py              |  2 ++
 dill/session.py            | 11 ++++++++---
 dill/tests/test_session.py |  2 +-
 3 files changed, 11 insertions(+), 4 deletions(-)

diff --git a/dill/_dill.py b/dill/_dill.py
index 787cb7ac..3ad7ead0 100644
--- a/dill/_dill.py
+++ b/dill/_dill.py
@@ -39,6 +39,7 @@
 from dill import logging
 from .logging import adapter as logger
 from .logging import trace as _trace
+_logger = logging.getLogger(__name__)
 
 import os
 import sys
@@ -1350,6 +1351,7 @@ def _save_module_dict(pickler, main_dict):
     is_builtin = _is_builtin_module(main)
     pickler.write(MARK + DICT)  # don't need to memoize
     for name, value in main_dict.items():
+        _logger.debug("Pickling %r (%s)", name, type(value).__name__)
         pickler.save(name)
         try:
             if pickler.save(value):
diff --git a/dill/session.py b/dill/session.py
index 9b52393c..9fc4ef56 100644
--- a/dill/session.py
+++ b/dill/session.py
@@ -87,9 +87,9 @@
 
 BUILTIN_CONSTANTS = (None, False, True, NotImplemented)
 
-def _stash_modules(main_module):
+def _stash_modules(main_module, original_main):
     """pop imported variables to be saved by reference in the __dill_imported* attributes"""
-    modmap = _module_map(main_module)
+    modmap = _module_map(original_main)
     newmod = ModuleType(main_module.__name__)
     original = {}
     imported = []
@@ -136,6 +136,7 @@ def _stash_modules(main_module):
             refimported += [(name, mod) for mod, name in imported_top_level]
             message = "[dump_module] Variables saved by reference (refimported):\n"
             logger.info(message + _format_log_dict(dict(refimported)))
+        logger.debug("main namespace after _stash_modules(): %s", dir(newmod))
 
         return newmod, modmap
     else:
@@ -176,6 +177,7 @@ def _filter_vars(main_module, exclude, include, base_rules):
     newmod = ModuleType(main_module.__name__)
     newmod.__dict__.update(namespace)
     _discard_added_variables(newmod, namespace)
+    logger.debug("main namespace after _filter_vars(): %s", dir(newmod))
     return newmod
 
 def _discard_added_variables(main, original_namespace):
@@ -363,9 +365,12 @@ def dump_module(
     if not isinstance(main, ModuleType):
         raise TypeError("%r is not a module" % main)
     original_main = main
+
+    logger.debug("original main namespace: %s", dir(main))
     main = _filter_vars(main, exclude, include, base_rules)
     if refimported:
-        main, modmap = _stash_modules(original_main)
+        main, modmap = _stash_modules(main, original_main)
+
     with _open(filename, 'wb', seekable=True) as file:
         pickler = Pickler(file, protocol, **kwds)
         pickler._main = main     #FIXME: dill.settings are disabled
diff --git a/dill/tests/test_session.py b/dill/tests/test_session.py
index de3e97af..e5341b25 100644
--- a/dill/tests/test_session.py
+++ b/dill/tests/test_session.py
@@ -201,7 +201,7 @@ def test_runtime_module():
     runtime_mod = ModuleType(modname)
     runtime_mod.x = 42
 
-    mod, _ = dill.session._stash_modules(runtime_mod)
+    mod, _ = dill.session._stash_modules(runtime_mod, runtime_mod)
     if mod is not runtime_mod:
         print("There are objects to save by referenece that shouldn't be:",
               mod.__dill_imported, mod.__dill_imported_as, mod.__dill_imported_top_level,