working on ipyflow integration

Author: smacke

core/superduperreload/patching.py

@@ -79,7 +79,7 @@ class ObjectPatcher:
 
     def __init__(self, patch_referrers: bool) -> None:
         self._patched_obj_ids: Set[int] = set()
-        self._patch_rules = [
+        self._patch_rules: List[Tuple[Callable, Callable]] = [
             (lambda a, b: isinstance2(a, b, type), self._patch_class),
             (lambda a, b: isinstance2(a, b, FunctionType), self._patch_function),
             (lambda a, b: isinstance2(a, b, MethodType), self._patch_method),
@@ -184,25 +184,31 @@ def _try_patch_readonly_attr(
             offset = cls._infer_field_offset(struct_type, new, field)
         cls._try_write_readonly_attr(struct_type, old, field, new_value, offset=offset)
 
-    def _patch_function(self, old, new):
+    def _patch_function(self, old: FunctionType, new: FunctionType) -> None:
         if old is new:
             return
         for name in _FUNC_ATTRS:
+            if name == "__globals__":
+                # switch order for __globals__ since we keep the old module.__dict__
+                old, new = new, old
             try:
                 setattr(old, name, getattr(new, name))
             except (AttributeError, TypeError, ValueError):
                 self._try_patch_readonly_attr(
                     _CPythonStructType.FUNCTION, old, new, name
                 )
+            finally:
+                if name == "__globals__":
+                    old, new = new, old
 
-    def _patch_method(self, old: MethodType, new: MethodType):
+    def _patch_method(self, old: MethodType, new: MethodType) -> None:
         if old is new:
             return
         self._patch_function(old.__func__, new.__func__)
         self._try_patch_readonly_attr(_CPythonStructType.METHOD, old, new, "__self__")
 
     @classmethod
-    def _patch_instances(cls, old, new):
+    def _patch_instances(cls, old: Type[object], new: Type[object]) -> None:
         """Use garbage collector to find all instances that refer to the old
         class definition and update their __class__ to point to the new class
         definition"""
@@ -269,7 +275,7 @@ def _patch_property(self, old: property, new: property) -> None:
     def _patch_partial(self, old: functools.partial, new: functools.partial) -> None:
         if old is new:
             return
-        self._patch_function(old.func, new.func)
+        self._patch_function(old.func, new.func)  # type: ignore
         self._try_patch_readonly_attr(_CPythonStructType.PARTIAL, old, new, "args")
         self._try_patch_readonly_attr(_CPythonStructType.PARTIAL, old, new, "keywords")
 

core/superduperreload/superduperreload.py

@@ -162,39 +162,41 @@ def filename_and_mtime(
 
         return py_filename, pymtime
 
-    def check(self, do_reload: bool = True) -> None:
-        """Check whether some modules need to be reloaded."""
-        self.reloaded_modules.clear()
-        self.failed_modules.clear()
-
-        # TODO: we should try to reload the modules in topological order
+    def _get_modules_needing_reload(self) -> Dict[str, Tuple[ModuleType, str, float]]:
+        modules_needing_reload = {}
         for modname, m in list(sys.modules.items()):
             package_components = modname.split(".")
             if any(
                 ".".join(package_components[:idx]) in self.skip_modules
                 for idx in range(1, len(package_components))
             ):
                 continue
-
             py_filename, pymtime = self.filename_and_mtime(m)
             if py_filename is None:
                 continue
-
-            try:
-                if pymtime <= self.modules_mtimes[modname]:
-                    continue
-            except KeyError:
-                self.modules_mtimes[modname] = pymtime
+            if pymtime <= self.modules_mtimes.setdefault(modname, pymtime):
+                continue
+            if self.failed.get(py_filename) == pymtime:
                 continue
-            else:
-                if self.failed.get(py_filename, None) == pymtime:
-                    continue
-
             self.modules_mtimes[modname] = pymtime
+            modules_needing_reload[modname] = (m, py_filename, pymtime)
+        return modules_needing_reload
 
-            if not do_reload:
-                continue
+    def check(self, do_reload: bool = True) -> None:
+        """Check whether some modules need to be reloaded."""
+        self.reloaded_modules.clear()
+        self.failed_modules.clear()
 
+        modules_needing_reload = self._get_modules_needing_reload()
+        if not do_reload:
+            return
+
+        # TODO: we should try to reload the modules in topological order
+        for modname, (
+            m,
+            py_filename,
+            pymtime,
+        ) in modules_needing_reload.items():
             # If we've reached this point, we should try to reload the module
             self._report(f"Reloading '{modname}'.")
             try:
@@ -220,6 +222,18 @@ def maybe_track_obj(self, module: ModuleType, name: str, obj: object) -> None:
         except TypeError:
             pass
 
+    def _patch_ipyflow_symbols(self, old, new, flow_):
+        if flow_ is None:
+            return
+        if isinstance(old, IMMUTABLE_PRIMITIVE_TYPES):
+            return
+        old_id = id(old)
+        if old_id not in flow_.aliases:
+            return
+        for sym in list(flow_.aliases[old_id]):
+            sym._override_ready_liveness_cell_num = flow_.cell_counter()
+            sym.update_obj_ref(new)
+
     def superduperreload(self, module: ModuleType) -> ModuleType:
         """Enhanced version of the superreload function from IPython's autoreload extension.
 
@@ -229,6 +243,13 @@ def superduperreload(self, module: ModuleType) -> ModuleType:
         - upgrades the code object of every old function and method
         - clears the module's namespace before reloading
         """
+        try:
+            from ipyflow import flow
+
+            flow_ = flow()
+        except:
+            flow_ = None
+
         self._patched_obj_ids.clear()
 
         # collect old objects in the module
@@ -254,6 +275,7 @@ def superduperreload(self, module: ModuleType) -> ModuleType:
                     continue
                 self._patch_generic(old_obj, new_obj)
                 self._patch_referrers_generic(old_obj, new_obj)
+                self._patch_ipyflow_symbols(old_obj, new_obj, flow_)
 
             if new_refs:
                 self.old_objects[key] = new_refs