Session: check id against module being saved

Author: leogama

dill/session.py

@@ -17,7 +17,7 @@
 
 import dill
 from dill import Pickler, Unpickler
-from ._dill import ModuleType, _import_module, _is_builtin_module
+from ._dill import ModuleType, _import_module, _is_builtin_module, _main_module
 from .utils import AttrDict, CheckerSet, TransSet
 from .settings import settings
 
@@ -111,10 +111,14 @@ def _exclude_objs(main, exclude_extra, filters_extra, settings):
     categories = {'ids': int, 'names': str, 'regex': re.Pattern, 'types': type}
     exclude = AttrDict({cat: copy(settings.session_exclude[cat]) for cat in categories})
     filters = copy(settings.session_filters)
+    del categories['ids']  # special case
     if exclude_extra is not None:
         if isinstance(exclude_extra, str):
             raise ValueError("'exclude' can be of type Iterable[str], but not str")
         for item in exclude_extra:
+            if isinstance(item, int):
+                exclude.ids.add(item, main=main)
+                continue
             for category, klass in categories.items():
                 if isinstance(item, klass):
                     exclude[category].add(item)
@@ -220,12 +224,12 @@ def load_session(filename: Union[os.PathLike, io.BytesIO] = '/tmp/session.pkl',
 #  Settings  #
 ##############
 
-def _as_id(item):
+def _as_id(item, *, main=_main_module):
     if isinstance(item, int):
-        import warnings, __main__
-        if not any(id(obj) == item for obj in __main__.__dict__.values()):
-            warnings.warn("%d isn't the id of any object in __main__ namespace. "
-                          "Did you mean 'id(%d)?'" % (item, item))
+        import warnings
+        if not any(id(obj) == item for obj in main.__dict__.values()):
+            warnings.warn("%d isn't the id of any object in the '%s' namespace. "
+                          "Did you mean 'id(%d)'?" % (item, main.__name__, item))
         return item
     return id(item)
 

dill/utils.py

@@ -47,8 +47,8 @@ class TransSet(set):
     def __init__(self, func: Callable, *args):
         self.constructor = func
         super().__init__(*args)
-    def add(self, item):
-        super().add(self.constructor(item))
+    def add(self, item, **kwargs):
+        super().add(self.constructor(item, **kwargs))
     def discard(self, item):
         super().discard(self.constructor(item))
     def remove(self, item):