| from __future__ import annotations |
| import ast |
| import builtins |
| import contextlib |
| import itertools |
| import logging |
| import os |
| import pickle |
| import platform |
| import sys |
| import textwrap |
| from types import ModuleType |
| from typing import TYPE_CHECKING, Any, NamedTuple, cast |
| from collections.abc import Generator, Iterable |
|
|
| from IPython.extensions.deduperreload.deduperreload_patching import ( |
| DeduperReloaderPatchingMixin, |
| ) |
|
|
| if TYPE_CHECKING: |
| TDefinitionAst = ( |
| ast.FunctionDef |
| | ast.AsyncFunctionDef |
| | ast.Import |
| | ast.ImportFrom |
| | ast.Assign |
| | ast.AnnAssign |
| ) |
|
|
|
|
| def get_module_file_name(module: ModuleType | str) -> str | None: |
| """Returns the module's file path, or the empty string if it's inaccessible""" |
| if (mod := sys.modules.get(module) if isinstance(module, str) else module) is None: |
| return "" |
| return getattr(mod, "__file__", "") or "" |
|
|
|
|
| def compare_ast(node1: ast.AST | list[ast.AST], node2: ast.AST | list[ast.AST]) -> bool: |
| """Checks if node1 and node2 have identical AST structure/values, apart from some attributes""" |
| if type(node1) is not type(node2): |
| return False |
|
|
| if isinstance(node1, ast.AST): |
| for k, v in node1.__dict__.items(): |
| if k in ( |
| "lineno", |
| "end_lineno", |
| "col_offset", |
| "end_col_offset", |
| "ctx", |
| "parent", |
| ): |
| continue |
| if not hasattr(node2, k) or not compare_ast(v, getattr(node2, k)): |
| return False |
| return True |
|
|
| elif isinstance(node1, list) and isinstance( |
| node2, list |
| ): |
| return len(node1) == len(node2) and all( |
| compare_ast(n1, n2) for n1, n2 in zip(node1, node2) |
| ) |
| else: |
| return node1 == node2 |
|
|
|
|
| class DependencyNode(NamedTuple): |
| """ |
| Each node represents a function. |
| qualified_name: string which represents the namespace/name of the function |
| abstract_syntax_tree: subtree of the overall module which corresponds to this function |
| |
| qualified_name is of the structure: (namespace1, namespace2, ..., name) |
| |
| For example, foo() in the following would be represented as (A, B, foo): |
| |
| class A: |
| class B: |
| def foo(): |
| pass |
| """ |
|
|
| qualified_name: tuple[str, ...] |
| abstract_syntax_tree: ast.AST |
|
|
|
|
| class GatherResult(NamedTuple): |
| import_defs: list[tuple[tuple[str, ...], ast.Import | ast.ImportFrom]] = [] |
| assign_defs: list[tuple[tuple[str, ...], ast.Assign | ast.AnnAssign]] = [] |
| function_defs: list[ |
| tuple[tuple[str, ...], ast.FunctionDef | ast.AsyncFunctionDef] |
| ] = [] |
| classes: dict[str, ast.ClassDef] = {} |
| unfixable: list[ast.AST] = [] |
|
|
| @classmethod |
| def create(cls) -> GatherResult: |
| return cls([], [], [], {}, []) |
|
|
| def all_defs(self) -> Iterable[tuple[tuple[str, ...], TDefinitionAst]]: |
| return itertools.chain(self.import_defs, self.assign_defs, self.function_defs) |
|
|
| def inplace_merge(self, other: GatherResult) -> None: |
| self.import_defs.extend(other.import_defs) |
| self.assign_defs.extend(other.assign_defs) |
| self.function_defs.extend(other.function_defs) |
| self.classes.update(other.classes) |
| self.unfixable.extend(other.unfixable) |
|
|
|
|
| class ConstexprDetector(ast.NodeVisitor): |
| def __init__(self) -> None: |
| self.is_constexpr = True |
| self._allow_builtins_exceptions = True |
|
|
| @contextlib.contextmanager |
| def disallow_builtins_exceptions(self) -> Generator[None, None, None]: |
| prev_allow = self._allow_builtins_exceptions |
| self._allow_builtins_exceptions = False |
| try: |
| yield |
| finally: |
| self._allow_builtins_exceptions = prev_allow |
|
|
| def visit_Attribute(self, node: ast.Attribute) -> None: |
| with self.disallow_builtins_exceptions(): |
| self.visit(node.value) |
|
|
| def visit_Name(self, node: ast.Name) -> None: |
| if self._allow_builtins_exceptions and hasattr(builtins, node.id): |
| return |
| self.is_constexpr = False |
|
|
| def visit(self, node: ast.AST) -> None: |
| if not self.is_constexpr: |
| |
| return |
| super().visit(node) |
|
|
| def __call__(self, node: ast.AST) -> bool: |
| self.is_constexpr = True |
| self.visit(node) |
| return self.is_constexpr |
|
|
|
|
| class AutoreloadTree: |
| """ |
| Recursive data structure to keep track of reloadable functions/methods. Each object corresponds to a specific scope level. |
| children: classes inside given scope, maps class name to autoreload tree for that class's scope |
| funcs_to_autoreload: list of function names that can be autoreloaded in given scope. |
| new_nested_classes: Classes getting added in new autoreload cycle |
| """ |
|
|
| def __init__(self) -> None: |
| self.children: dict[str, AutoreloadTree] = {} |
| self.defs_to_reload: list[tuple[tuple[str, ...], ast.AST]] = [] |
| self.defs_to_delete: set[str] = set() |
| self.new_nested_classes: dict[str, ast.AST] = {} |
|
|
| def traverse_prefixes(self, prefixes: list[str]) -> AutoreloadTree: |
| """ |
| Return ref to the AutoreloadTree at the namespace specified by prefixes |
| """ |
| cur = self |
| for prefix in prefixes: |
| if prefix not in cur.children: |
| cur.children[prefix] = AutoreloadTree() |
| cur = cur.children[prefix] |
| return cur |
|
|
|
|
| class DeduperReloader(DeduperReloaderPatchingMixin): |
| """ |
| This version of autoreload detects when we can leverage targeted recompilation of a subset of a module and patching |
| existing function/method objects to reflect these changes. |
| |
| Detects what functions/methods can be reloaded by recursively comparing the old/new AST of module-level classes, |
| module-level classes' methods, recursing through nested classes' methods. If other changes are made, original |
| autoreload algorithm is called directly. |
| """ |
|
|
| def __init__(self) -> None: |
| self._to_autoreload: AutoreloadTree = AutoreloadTree() |
| self.source_by_modname: dict[str, str] = {} |
| self.dependency_graph: dict[tuple[str, ...], list[DependencyNode]] = {} |
| self._enabled = True |
|
|
| @property |
| def enabled(self) -> bool: |
| return self._enabled and platform.python_implementation() == "CPython" |
|
|
| @enabled.setter |
| def enabled(self, value: bool) -> None: |
| self._enabled = value |
|
|
| def update_sources(self) -> None: |
| """ |
| Update dictionary source_by_modname with current modules' source codes. |
| """ |
| if not self.enabled: |
| return |
| for new_modname in sys.modules.keys() - self.source_by_modname.keys(): |
| new_module = sys.modules[new_modname] |
| if ( |
| (fname := get_module_file_name(new_module)) is None |
| or "site-packages" in fname |
| or "dist-packages" in fname |
| or not os.access(fname, os.R_OK) |
| ): |
| self.source_by_modname[new_modname] = "" |
| continue |
| |
| if not fname.endswith(".py"): |
| self.source_by_modname[new_modname] = "" |
| continue |
| try: |
| with open(fname, "r") as f: |
| self.source_by_modname[new_modname] = f.read() |
| except Exception as e: |
| logger = logging.getLogger("autoreload") |
| logger.exception( |
| f"Failed to read module file '{fname}' for module '{new_modname}': {type(e).__name__}" |
| ) |
| self.source_by_modname[new_modname] = "" |
|
|
| constexpr_detector = ConstexprDetector() |
|
|
| @staticmethod |
| def is_enum_subclass(node: ast.Module | ast.ClassDef) -> bool: |
| if isinstance(node, ast.Module): |
| return False |
| for base in node.bases: |
| if isinstance(base, ast.Name) and base.id == "Enum": |
| return True |
| elif ( |
| isinstance(base, ast.Attribute) |
| and base.attr == "Enum" |
| and isinstance(base.value, ast.Name) |
| and base.value.id == "enum" |
| ): |
| return True |
| return False |
|
|
| @classmethod |
| def is_constexpr_assign( |
| cls, node: ast.AST, parent_node: ast.Module | ast.ClassDef |
| ) -> bool: |
| if not isinstance(node, (ast.Assign, ast.AnnAssign)) or node.value is None: |
| return False |
| if cls.is_enum_subclass(parent_node): |
| return False |
| for target in node.targets if isinstance(node, ast.Assign) else [node.target]: |
| if not isinstance(target, ast.Name): |
| return False |
| return cls.constexpr_detector(node.value) |
|
|
| @classmethod |
| def _gather_children( |
| cls, body: list[ast.stmt], parent_node: ast.Module | ast.ClassDef |
| ) -> GatherResult: |
| """ |
| Given list of ast elements, return: |
| 1. dict mapping function names to their ASTs. |
| 2. dict mapping class names to their ASTs. |
| 3. list of any other ASTs. |
| """ |
| result = GatherResult.create() |
| for ast_node in body: |
| ast_elt: ast.expr | ast.stmt = ast_node |
| while isinstance(ast_elt, ast.Expr): |
| ast_elt = ast_elt.value |
| if isinstance(ast_elt, (ast.FunctionDef, ast.AsyncFunctionDef)): |
| result.function_defs.append(((ast_elt.name,), ast_elt)) |
| elif isinstance(ast_elt, (ast.Import, ast.ImportFrom)): |
| result.import_defs.append( |
| (tuple(name.asname or name.name for name in ast_elt.names), ast_elt) |
| ) |
| elif isinstance(ast_elt, ast.ClassDef): |
| result.classes[ast_elt.name] = ast_elt |
| elif isinstance(ast_elt, ast.If): |
| result.unfixable.append(ast_elt.test) |
| result.inplace_merge(cls._gather_children(ast_elt.body, parent_node)) |
| result.inplace_merge(cls._gather_children(ast_elt.orelse, parent_node)) |
| elif isinstance(ast_elt, (ast.AsyncWith, ast.With)): |
| result.unfixable.extend(ast_elt.items) |
| result.inplace_merge(cls._gather_children(ast_elt.body, parent_node)) |
| elif isinstance(ast_elt, ast.Try): |
| result.inplace_merge(cls._gather_children(ast_elt.body, parent_node)) |
| result.inplace_merge(cls._gather_children(ast_elt.orelse, parent_node)) |
| result.inplace_merge( |
| cls._gather_children(ast_elt.finalbody, parent_node) |
| ) |
| for handler in ast_elt.handlers: |
| if handler.type is not None: |
| result.unfixable.append(handler.type) |
| result.inplace_merge( |
| cls._gather_children(handler.body, parent_node) |
| ) |
| elif not isinstance(ast_elt, (ast.Constant, ast.Pass)): |
| if cls.is_constexpr_assign(ast_elt, parent_node): |
| assert isinstance(ast_elt, (ast.Assign, ast.AnnAssign)) |
| targets = ( |
| ast_elt.targets |
| if isinstance(ast_elt, ast.Assign) |
| else [ast_elt.target] |
| ) |
| result.assign_defs.append( |
| ( |
| tuple(cast(ast.Name, target).id for target in targets), |
| ast_elt, |
| ) |
| ) |
| else: |
| result.unfixable.append(ast_elt) |
| return result |
|
|
| def detect_autoreload( |
| self, |
| old_node: ast.Module | ast.ClassDef, |
| new_node: ast.Module | ast.ClassDef, |
| prefixes: list[str] | None = None, |
| ) -> bool: |
| """ |
| Returns |
| ------- |
| `True` if we can run our targeted autoreload algorithm safely. |
| `False` if we should instead use IPython's original autoreload implementation. |
| """ |
| if not self.enabled: |
| return False |
| prefixes = prefixes or [] |
|
|
| old_result = self._gather_children(old_node.body, old_node) |
| new_result = self._gather_children(new_node.body, new_node) |
| old_defs_by_name: dict[str, ast.AST] = { |
| name: ast_def for names, ast_def in old_result.all_defs() for name in names |
| } |
| new_defs_by_name: dict[str, ast.AST] = { |
| name: ast_def for names, ast_def in new_result.all_defs() for name in names |
| } |
|
|
| if not compare_ast(old_result.unfixable, new_result.unfixable): |
| return False |
|
|
| cur = self._to_autoreload.traverse_prefixes(prefixes) |
| for names, new_ast_def in new_result.all_defs(): |
| names_to_reload = [] |
| for name in names: |
| if new_defs_by_name[name] is not new_ast_def: |
| continue |
| if name not in old_defs_by_name or not compare_ast( |
| new_ast_def, old_defs_by_name[name] |
| ): |
| names_to_reload.append(name) |
| if names_to_reload: |
| cur.defs_to_reload.append((tuple(names), new_ast_def)) |
| cur.defs_to_delete |= set(old_defs_by_name.keys()) - set( |
| new_defs_by_name.keys() |
| ) |
| for name, new_ast_def_class in new_result.classes.items(): |
| if name not in old_result.classes: |
| cur.new_nested_classes[name] = new_ast_def_class |
| elif not compare_ast( |
| new_ast_def_class, old_result.classes[name] |
| ) and not self.detect_autoreload( |
| old_result.classes[name], new_ast_def_class, prefixes + [name] |
| ): |
| return False |
| return True |
|
|
| def _check_dependents(self) -> bool: |
| """ |
| If a decorator function is modified, we should similarly reload the functions which are decorated by this |
| decorator. Iterate through the Dependency Graph to find such cases in the given AutoreloadTree. |
| """ |
| for node in self._check_dependents_inner(): |
| self._add_node_to_autoreload_tree(node) |
| return True |
|
|
| def _add_node_to_autoreload_tree(self, node: DependencyNode) -> None: |
| """ |
| Given a node of the dependency graph, add decorator dependencies to the autoreload tree. |
| """ |
| if len(node.qualified_name) == 0: |
| return |
| cur = self._to_autoreload.traverse_prefixes(list(node.qualified_name[:-1])) |
| if node.abstract_syntax_tree is not None: |
| cur.defs_to_reload.append( |
| ((node.qualified_name[-1],), node.abstract_syntax_tree) |
| ) |
|
|
| def _check_dependents_inner( |
| self, prefixes: list[str] | None = None |
| ) -> list[DependencyNode]: |
| prefixes = prefixes or [] |
| cur = self._to_autoreload.traverse_prefixes(prefixes) |
| ans = [] |
| for (func_name, *_), _ in cur.defs_to_reload: |
| node = tuple(prefixes + [func_name]) |
| ans.extend(self._gen_dependents(node)) |
| for class_name in cur.new_nested_classes: |
| ans.extend(self._check_dependents_inner(prefixes + [class_name])) |
| return ans |
|
|
| def _gen_dependents(self, qualname: tuple[str, ...]) -> list[DependencyNode]: |
| ans = [] |
| if qualname not in self.dependency_graph: |
| return [] |
| for elt in self.dependency_graph[qualname]: |
| ans.extend(self._gen_dependents(elt.qualified_name)) |
| ans.append(elt) |
| return ans |
|
|
| def _patch_namespace_inner( |
| self, ns: ModuleType | type, prefixes: list[str] | None = None |
| ) -> bool: |
| """ |
| This function patches module functions and methods. Specifically, only objects with their name in |
| self.to_autoreload will be considered for patching. If an object has been marked to be autoreloaded, |
| new_source_code gets executed in the old version's global environment. Then, replace the old function's |
| attributes with the new function's attributes. |
| """ |
| prefixes = prefixes or [] |
| cur = self._to_autoreload.traverse_prefixes(prefixes) |
| namespace_to_check = ns |
| for prefix in prefixes: |
| namespace_to_check = namespace_to_check.__dict__[prefix] |
| seen_names: set[str] = set() |
| for names, new_ast_def in cur.defs_to_reload: |
| if len(names) == 1 and names[0] in seen_names: |
| continue |
| seen_names.update(names) |
| local_env: dict[str, Any] = {} |
| if ( |
| isinstance(new_ast_def, (ast.FunctionDef, ast.AsyncFunctionDef)) |
| and (name := names[0]) in namespace_to_check.__dict__ |
| ): |
| assert len(names) == 1 |
| to_patch_to = namespace_to_check.__dict__[name] |
| if isinstance(to_patch_to, (staticmethod, classmethod)): |
| to_patch_to = to_patch_to.__func__ |
| |
| func_code = textwrap.dedent(ast.unparse(new_ast_def)) |
| if is_method := (len(prefixes) > 0): |
| func_code = "class __autoreload_class__:\n" + textwrap.indent( |
| func_code, " " |
| ) |
| global_env = ns.__dict__ |
| if not isinstance(global_env, dict): |
| global_env = dict(global_env) |
| |
| filename = ( |
| getattr(to_patch_to, "__code__", None) |
| and to_patch_to.__code__.co_filename |
| or "<string>" |
| ) |
| func_asts = [ast.parse(func_code)] |
| if len(cast(ast.FunctionDef, func_asts[0].body[0]).decorator_list) > 0: |
| without_decorator_list = pickle.loads(pickle.dumps(func_asts[0])) |
| cast( |
| ast.FunctionDef, without_decorator_list.body[0] |
| ).decorator_list = [] |
| func_asts.insert(0, without_decorator_list) |
| for func_ast in func_asts: |
| compiled_code = compile( |
| func_ast, filename, mode="exec", dont_inherit=True |
| ) |
| exec(compiled_code, global_env, local_env) |
| |
| if is_method: |
| to_patch_from = getattr(local_env["__autoreload_class__"], name) |
| else: |
| to_patch_from = local_env[name] |
| if isinstance(to_patch_from, (staticmethod, classmethod)): |
| to_patch_from = to_patch_from.__func__ |
| if isinstance(to_patch_to, property) and isinstance( |
| to_patch_from, property |
| ): |
| for attr in ("fget", "fset", "fdel"): |
| if ( |
| getattr(to_patch_to, attr) is None |
| or getattr(to_patch_from, attr) is None |
| ): |
| self.try_patch_attr(to_patch_to, to_patch_from, attr) |
| else: |
| self.patch_function( |
| getattr(to_patch_to, attr), |
| getattr(to_patch_from, attr), |
| is_method, |
| ) |
| elif not isinstance(to_patch_to, property) and not isinstance( |
| to_patch_from, property |
| ): |
| self.patch_function(to_patch_to, to_patch_from, is_method) |
| else: |
| raise ValueError( |
| "adding or removing property decorations not supported" |
| ) |
| else: |
| exec( |
| ast.unparse(new_ast_def), |
| ns.__dict__ | namespace_to_check.__dict__, |
| local_env, |
| ) |
| for name in names: |
| setattr(namespace_to_check, name, local_env[name]) |
| cur.defs_to_reload.clear() |
| for name in cur.defs_to_delete: |
| try: |
| delattr(namespace_to_check, name) |
| except (AttributeError, TypeError, ValueError): |
| |
| pass |
| cur.defs_to_delete.clear() |
| for class_name, class_ast_node in cur.new_nested_classes.items(): |
| local_env_class: dict[str, Any] = {} |
| exec( |
| ast.unparse(class_ast_node), |
| ns.__dict__ | namespace_to_check.__dict__, |
| local_env_class, |
| ) |
| setattr(namespace_to_check, class_name, local_env_class[class_name]) |
| cur.new_nested_classes.clear() |
| for class_name in cur.children.keys(): |
| if not self._patch_namespace(ns, prefixes + [class_name]): |
| return False |
| cur.children.clear() |
| return True |
|
|
| def _patch_namespace( |
| self, ns: ModuleType | type, prefixes: list[str] | None = None |
| ) -> bool: |
| """ |
| Wrapper for patching all elements in a namespace as specified by the to_autoreload member variable. |
| Returns `true` if patching was successful, and `false` if unsuccessful. |
| """ |
| try: |
| return self._patch_namespace_inner(ns, prefixes=prefixes) |
| except Exception: |
| return False |
|
|
| def maybe_reload_module(self, module: ModuleType) -> bool: |
| """ |
| Uses Deduperreload to try to update a module. |
| Returns `true` on success and `false` on failure. |
| """ |
| if not self.enabled: |
| return False |
| if not (modname := getattr(module, "__name__", None)): |
| return False |
| if (fname := get_module_file_name(module)) is None: |
| return False |
| try: |
| with open(fname, "r") as f: |
| new_source_code = f.read() |
| except Exception as e: |
| logger = logging.getLogger("autoreload") |
| logger.exception( |
| f"Failed to read module file '{fname}' for module '{modname}': {type(e).__name__}" |
| ) |
| return False |
| patched_flag = False |
| if old_source_code := self.source_by_modname.get(modname): |
| |
| try: |
| old_module_ast = ast.parse(old_source_code) |
| new_module_ast = ast.parse(new_source_code) |
| except Exception: |
| return False |
| |
| ctx = contextlib.suppress() |
| with ctx: |
| self._build_dependency_graph(new_module_ast) |
| if ( |
| self.detect_autoreload(old_module_ast, new_module_ast) |
| and self._check_dependents() |
| and self._patch_namespace(module) |
| ): |
| patched_flag = True |
|
|
| self.source_by_modname[modname] = new_source_code |
| self._to_autoreload = AutoreloadTree() |
| return patched_flag |
|
|
| def _separate_name( |
| self, |
| decorator: ast.Attribute | ast.Name | ast.Call | ast.expr, |
| accept_calls: bool, |
| ) -> list[str] | None: |
| """ |
| Generates a qualified name for a given decorator by finding its relative namespace. |
| """ |
| if isinstance(decorator, ast.Name): |
| return [decorator.id] |
| elif isinstance(decorator, ast.Call): |
| if accept_calls: |
| return self._separate_name(decorator.func, False) |
| else: |
| return None |
| if not isinstance(decorator, ast.Attribute): |
| return None |
| if pref := self._separate_name(decorator.value, False): |
| return pref + [decorator.attr] |
| else: |
| return None |
|
|
| def _gather_dependents( |
| self, body: list[ast.stmt], body_prefixes: list[str] | None = None |
| ) -> bool: |
| body_prefixes = body_prefixes or [] |
| for ast_node in body: |
| ast_elt: ast.expr | ast.stmt = ast_node |
| if isinstance(ast_elt, ast.ClassDef): |
| self._gather_dependents(ast_elt.body, body_prefixes + [ast_elt.name]) |
| continue |
| if not isinstance(ast_elt, (ast.FunctionDef, ast.AsyncFunctionDef)): |
| continue |
| qualified_name = tuple(body_prefixes + [ast_elt.name]) |
| cur_dependency_node = DependencyNode(qualified_name, ast_elt) |
| for decorator in ast_elt.decorator_list: |
| decorator_path = self._separate_name(decorator, True) |
| if not decorator_path: |
| continue |
| decorator_path_tuple = tuple(decorator_path) |
| self.dependency_graph.setdefault(decorator_path_tuple, []).append( |
| cur_dependency_node |
| ) |
| return True |
|
|
| def _build_dependency_graph(self, new_ast: ast.Module | ast.ClassDef) -> bool: |
| """ |
| Wrapper function for generating dependency graph given some AST. |
| Returns `true` on success. Returns `false` on failure. |
| Currently, only returns `true` as we do not block on failure to build this graph. |
| """ |
| return self._gather_dependents(new_ast.body) |
|
|
