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( # type:ignore [redundant-expr]

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:

# can short-circuit if we've already detected that it's not a 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

# Skip binary files (e.g., .so, .pyd, .pyo)

if not fname.endswith(".py"):

self.source_by_modname[new_modname] = ""

continue

try:

with open(fname, "r", encoding="utf8") 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__

# exec new source code using old function's (obj) globals environment.

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)

# Compile with correct filename to preserve in traceback

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) # type: ignore[arg-type]

# local_env contains the function exec'd from new version of function

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):

# give up on deleting the attribute, let the stale one dangle

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", encoding="utf8") 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):

# get old/new module ast

try:

old_module_ast = ast.parse(old_source_code)

new_module_ast = ast.parse(new_source_code)

except Exception:

return False

# detect if we are able to use our autoreload algorithm

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)