#!/usr/bin/env python
from __future__ import annotations
import ast
import os
import platform
import pytest
import random
import shutil
import sys
import tempfile
import textwrap
import time
from types import ModuleType

from IPython.extensions.autoreload import AutoreloadMagics

from IPython.extensions.deduperreload.deduperreload import compare_ast, DeduperReloader

if platform.python_implementation() != "CPython":
    pytest.skip(
        "We do not support non-CPython versions.",
        allow_module_level=True,
    )


class DeduperTestReloader(DeduperReloader):
    def __init__(self, *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)
        self.exceptions_raised: list[Exception] = []

    def _patch_namespace(
        self, module: ModuleType | type, prefixes: list[str] | None = None
    ) -> bool:
        try:
            assert super()._patch_namespace(module, prefixes)
            return True
        except Exception as e:
            self.exceptions_raised.append(e)
            return False


def is_nonempty_file(fname: str) -> bool:
    return len(fname) > 0


def squish_text(text: str) -> str:
    """
    Turns text like this:

    '''        def foo():
    return "bar"
            def baz():
                return "bat"
    def bam():
                return "bat"
    '''

    into this:

    '''def foo():
        return "bar"
    def baz():
        return "bat"
    def bam():
        return "bat"
    '''

    The former is common when we are trying to use string templates
    whose parameters are multiline and unaware of the existing indentation.

    :param text: a string with messed up indentation
    :return: `text` but with indentation fixed
    """
    prev_indentation = 0
    transformed_text_lines = []
    for line in text.strip("\n").splitlines():
        line_without_indentation = line.lstrip()
        indentation = len(line) - len(line_without_indentation)
        if indentation == 0:
            indentation = prev_indentation
        else:
            prev_indentation = indentation
        transformed_text_lines.append(
            textwrap.indent(line_without_indentation, " " * indentation)
        )
    return textwrap.dedent("\n".join(transformed_text_lines))


# ---------------------------------------------------------------------------
# AutoreloadDetectionSuite
# ---------------------------------------------------------------------------


def test_compare_ast():
    code1 = squish_text(
        """
        def factorial(n):
            def fn(sdfsdf):
                print(sdfsdf)
            if n < 0:
                return "Factorial is not defined for negative numbers."
            elif n == 0 or n == 1:
                return 1
            else:
                result = 1
                for i in range(2, n + 1):
                    result *= i
                return result
            y = 12
            print(y)
        print(1)
        x = 1212121
    """
    )
    code2 = squish_text(
        """
        def factorial(n):
            def fn(sdfsdf):
                print(sdfsdf+"!!!")
            if n < 0:
                return "Factorial is not defined for negative numbers."
            elif n == 0 or n == 1:
                return 1
            else:
                result = 1
                for i in range(2, n + 1):
                    result *= i
                return result
            y = 12
            print(y)
        print(1)
        x = 1212121
    """
    )
    ast_1 = ast.parse(code1)
    ast_2 = ast.parse(code2)

    assert not compare_ast(ast_1, ast_2)


def test_compare_ast2():
    code1 = squish_text(
        """
        def factorial(n):
            def fn(sdfsdf):
                print(sdfsdf)
            if n < 0:
                return "Factorial is not defined for negative numbers."
            elif n == 0 or n == 1:
                return 1
            else:
                result = 1
                for i in range(2, n + 1):
                    result *= i
                return result
            y = 12
            print(y)
        print(1)
        x = 1212121
    """
    )
    code2 = squish_text(
        """
        def factorial(n):
            def fn(sdfsdf):
                print(sdfsdf)
            if n < 0:
                return "Factorial is not defined for negative numbers."
            elif n == 0 or n == 1:
                return 1
            else:
                result = 1
                for i in range(2, n + 1):
                    result *= i
                return result
            y = 12
            print(y)
        print(1)
        x = 1212121
    """
    )
    ast_1 = ast.parse(code1)
    ast_2 = ast.parse(code2)

    assert compare_ast(ast_1, ast_2)


def test_autoreload_no_changes():
    code1 = squish_text(
        """
        def factorial(n):
            print(n)
        print(1)
        x = 1212121
    """
    )
    code2 = squish_text(
        """
        def factorial(n):
            print(n)
        print(1)
        x = 1212121
    """
    )
    deduperreloader = DeduperTestReloader()
    ast_1 = ast.parse(code1)
    ast_2 = ast.parse(code2)

    assert deduperreloader.detect_autoreload(ast_1, ast_2)
    assert deduperreloader._to_autoreload.defs_to_reload == []


def test_autoreload_static_assign_change_outside_function():
    code1 = squish_text(
        """
        def factorial(n):
            print(n)
        print(1)
    """
    )
    code2 = squish_text(
        """
        def factorial(n):
            print(n)
        print(1)
        x = 1212121
    """
    )
    deduperreloader = DeduperTestReloader()
    ast_1 = ast.parse(code1)
    ast_2 = ast.parse(code2)

    assert deduperreloader.detect_autoreload(ast_1, ast_2)


def test_autoreload_changes_inside_function():
    code1 = squish_text(
        """
        def factorial(n):
            print(n)
        print(1)
    """
    )
    code2 = squish_text(
        """
        def factorial(n):
            print(n + "edit!")
        print(1)
    """
    )
    deduperreloader = DeduperTestReloader()
    ast_1 = ast.parse(code1)
    ast_2 = ast.parse(code2)

    assert deduperreloader.detect_autoreload(ast_1, ast_2)
    assert len(deduperreloader._to_autoreload.defs_to_reload) == 1


def test_autoreload_changes_inside_and_outside_function():
    code1 = squish_text(
        """
        def factorial(n):
            print(n)
        print(1)
    """
    )
    code2 = squish_text(
        """
        def factorial(n):
            print(n + "edit!")
        print(2)
    """
    )
    deduperreloader = DeduperTestReloader()
    ast_1 = ast.parse(code1)
    ast_2 = ast.parse(code2)

    assert not deduperreloader.detect_autoreload(ast_1, ast_2)


def test_autoreload_changes_inner_function():
    code1 = squish_text(
        """
        def factorial(n):
            def foo():
                x = 3
                return x
            return foo
    """
    )
    code2 = squish_text(
        """
        def factorial(n):
            def foo():
                x = 4
                return x
            return foo
    """
    )
    deduperreloader = DeduperTestReloader()
    ast_1 = ast.parse(code1)
    ast_2 = ast.parse(code2)

    assert deduperreloader.detect_autoreload(ast_1, ast_2)
    assert len(deduperreloader._to_autoreload.defs_to_reload) == 1


def test_autoreload_changes_multiple_function():
    code1 = squish_text(
        """
        def factorial(n):
            def foo():
                x = 3
                return x
            return foo
        def bar():
            return 1
    """
    )
    code2 = squish_text(
        """
        def factorial(n):
            def foo():
                x = 4
                return x
            return foo
        def bar():
            return 2
    """
    )
    deduperreloader = DeduperTestReloader()
    ast_1 = ast.parse(code1)
    ast_2 = ast.parse(code2)

    assert deduperreloader.detect_autoreload(ast_1, ast_2)
    assert len(deduperreloader._to_autoreload.defs_to_reload) == 2


def test_autoreload_change_one_function_of_multiple():
    code1 = squish_text(
        """
        def factorial(n):
            def foo():
                x = 3
                return x
            return foo
        def bar():
            return 1
    """
    )
    code2 = squish_text(
        """
        def factorial(n):
            def foo():
                x = 4
                return x
            return foo
        def bar():
            return 1
    """
    )
    deduperreloader = DeduperTestReloader()
    ast_1 = ast.parse(code1)
    ast_2 = ast.parse(code2)

    assert deduperreloader.detect_autoreload(ast_1, ast_2)
    assert len(deduperreloader._to_autoreload.defs_to_reload) == 1


def test_autoreload_handling_moves():
    code1 = squish_text(
        """
        def factorial(n):
            return 1
        x = 1
    """
    )
    code2 = squish_text(
        """
        x = 1
        def factorial(n):
            return 1
    """
    )
    deduperreloader = DeduperTestReloader()
    ast_1 = ast.parse(code1)
    ast_2 = ast.parse(code2)

    assert deduperreloader.detect_autoreload(ast_1, ast_2)
    assert len(deduperreloader._to_autoreload.defs_to_reload) == 0


def test_autoreload_handling_function_moves_success():
    code1 = squish_text(
        """
        def factorial(n):
            return 1
        x = 1
        def foo():
            return 23
        x = 2
    """
    )
    code2 = squish_text(
        """
    x = 1
    x = 2
    def foo():
        return 23
    def factorial(n):
        return 1
    """
    )
    deduperreloader = DeduperTestReloader()
    ast_1 = ast.parse(code1)
    ast_2 = ast.parse(code2)

    assert deduperreloader.detect_autoreload(ast_1, ast_2)
    assert len(deduperreloader._to_autoreload.defs_to_reload) == 0


def test_autoreload_handling_function_moves_only():
    code1 = squish_text(
        """
    def factorial(n):
        return 1
    x = 1
    x = 2
    """
    )
    code2 = squish_text(
        """
    x = 1
    x = 2
    def factorial(n):
        return 1
    """
    )
    deduperreloader = DeduperTestReloader()
    ast_1 = ast.parse(code1)
    ast_2 = ast.parse(code2)

    assert deduperreloader.detect_autoreload(ast_1, ast_2)


def test_autoreload_handles_new_imports():
    code1 = squish_text(
        """
    def factorial(n):
        return 1
    x = 1
    """
    )
    code2 = squish_text(
        """
    import ast
    def factorial(n):
        return 1
    x = 1
    """
    )
    deduperreloader = DeduperTestReloader()
    ast_1 = ast.parse(code1)
    ast_2 = ast.parse(code2)

    assert deduperreloader.detect_autoreload(ast_1, ast_2)


def test_autoreload_async_function():
    code1 = squish_text(
        """
    async def sleep():
        print(f'Time: {time.time() - start:.2f}')
        await asyncio.sleep(1)
    """
    )
    code2 = squish_text(
        """
    async def sleep():
        print(f'Time: {time.time() - start:.2f}')
        await asyncio.sleep(10)
    """
    )
    deduperreloader = DeduperTestReloader()
    ast_1 = ast.parse(code1)
    ast_2 = ast.parse(code2)

    assert deduperreloader.detect_autoreload(ast_1, ast_2)
    assert len(deduperreloader._to_autoreload.defs_to_reload) == 1


def test_autoreload_add_function():
    code1 = squish_text(
        """
    async def sleep():
        print(f'Time: {time.time() - start:.2f}')
        await asyncio.sleep(1)
    """
    )
    code2 = squish_text(
        """
    async def sleep():
        print(f'Time: {time.time() - start:.2f}')
        await asyncio.sleep(1)
    def add(x,y):
        pass
    """
    )
    deduperreloader = DeduperTestReloader()
    ast_1 = ast.parse(code1)
    ast_2 = ast.parse(code2)

    assert deduperreloader.detect_autoreload(ast_1, ast_2)
    assert list(d[0][0] for d in deduperreloader._to_autoreload.defs_to_reload) == [
        "add"
    ]


def test_autoreload_add_function_ellipsis():
    code1 = squish_text(
        """
    async def sleep():
        print(f'Time: {time.time() - start:.2f}')
        await asyncio.sleep(1)
    """
    )
    code2 = squish_text(
        """
    async def sleep():
        print(f'Time: {time.time() - start:.2f}')
        await asyncio.sleep(1)
    def add(x,y):
        ...
    """
    )
    deduperreloader = DeduperTestReloader()
    ast_1 = ast.parse(code1)
    ast_2 = ast.parse(code2)

    assert deduperreloader.detect_autoreload(ast_1, ast_2)
    assert list(d[0][0] for d in deduperreloader._to_autoreload.defs_to_reload) == [
        "add"
    ]


# ---------------------------------------------------------------------------
# AutoreloadPatchingSuite — fixture
# ---------------------------------------------------------------------------


@pytest.fixture
def deduperreloader():
    return DeduperTestReloader()


def test_patching(deduperreloader):
    code1 = squish_text(
        """
        def foo():
            return 1
    """
    )
    code2 = squish_text(
        """
        def foo():
            return 2
    """
    )
    deduperreloader._to_autoreload.defs_to_reload = [
        (("foo",), ast.parse(code2))
    ]
    mod = ModuleType("mod")
    exec(code1, mod.__dict__)
    deduperreloader._patch_namespace(mod)
    assert mod.foo() == 2


def test_patching_parameters(deduperreloader):
    code1 = squish_text(
        """
        def foo(n,s):
            return n+s
    """
    )
    code2 = squish_text(
        """
        def foo(n):
            return n
    """
    )
    deduperreloader._to_autoreload.defs_to_reload = [
        (("foo",), ast.parse(code2))
    ]
    mod = ModuleType("mod")
    exec(code1, mod.__dict__)
    deduperreloader._patch_namespace(mod)
    assert mod.foo(2) == 2


def test_add_function(deduperreloader):
    code1 = squish_text(
        """
        def foo2(n):
            return n
    """
    )
    code2 = squish_text(
        """
        def foo(n):
            return 55
    """
    )
    deduperreloader._to_autoreload.defs_to_reload = [
        (("foo",), ast.parse(code2))
    ]
    mod = ModuleType("mod")
    exec(code1, mod.__dict__)
    deduperreloader._patch_namespace(mod)
    assert mod.foo(2) == 55
    assert mod.foo2(2) == 2


def test_two_operations(deduperreloader):
    code1 = squish_text(
        """
        def foo(n):
            return 1
    """
    )
    code2 = squish_text(
        """
        def foo(n):
            return 55
    """
    )
    code3 = squish_text(
        """
        def goo():
            return -1
        def foo(n):
            x = 2
            return x+n
        def bar():
            return 200
    """
    )
    deduperreloader._to_autoreload.defs_to_reload = [
        (("foo",), ast.parse(code2))
    ]
    mod = ModuleType("mod")
    exec(code1, mod.__dict__)
    assert mod.foo(2) == 1
    deduperreloader._patch_namespace(mod)
    assert mod.foo(2) == 55
    deduperreloader._to_autoreload.defs_to_reload = [
        (("foo",), ast.parse(code3))
    ]
    deduperreloader._patch_namespace(mod)
    assert mod.foo(2) == 4


def test_using_outside_param(deduperreloader):
    code1 = squish_text(
        """
        x=1
        def foo(n):
            return 1
    """
    )
    code2 = squish_text(
        """
        x=1
        def foo(n):
            return 55+x
    """
    )
    deduperreloader._to_autoreload.defs_to_reload = [
        (("foo",), ast.parse(code2))
    ]
    mod = ModuleType("mod")
    exec(code1, mod.__dict__)
    assert mod.foo(2) == 1
    deduperreloader._patch_namespace(mod)
    assert mod.foo(2) == 56


def test_importing_func(deduperreloader):
    code1 = squish_text(
        """
        from os import environ
        def foo(n):
            pass
    """
    )
    code2 = squish_text(
        """
        from os import environ
        def foo():
            environ._data
            return 1
    """
    )
    deduperreloader._to_autoreload.defs_to_reload = [
        (("foo",), ast.parse(code2))
    ]
    mod = ModuleType("mod")
    exec(code1, mod.__dict__)
    deduperreloader._patch_namespace(mod)
    assert mod.foo() == 1


# ---------------------------------------------------------------------------
# ShellFixture — helper class and pytest fixture
# ---------------------------------------------------------------------------


class FakeShell:
    def __init__(self):
        self.ns = {}
        self.user_ns = self.ns
        self.user_ns_hidden = {}
        self.auto_magics = AutoreloadMagics(shell=self)

    @staticmethod
    def pre_run_cell(obj):
        try_with_arg = False
        try:
            obj.pre_run_cell()
        except TypeError:
            try_with_arg = True
        if try_with_arg:
            obj.pre_run_cell(None)

    @staticmethod
    def post_run_cell(obj):
        try_with_arg = False
        try:
            obj.post_run_cell()
        except TypeError:
            try_with_arg = True
        if try_with_arg:
            obj.post_run_cell(None)

    def run_code(self, code):
        self.pre_run_cell(self.auto_magics)
        exec(code, self.user_ns)
        self.auto_magics.post_execute_hook()

    def push(self, items):
        self.ns.update(items)

    def magic_autoreload(self, parameter):
        self.auto_magics.autoreload(parameter)


class _ShellFixtureHelper:
    """Helper for creating test module files"""

    filename_chars = "abcdefghijklmopqrstuvwxyz0123456789"

    def __init__(self):
        self.created_temp_modules = set()
        self.test_dir = tempfile.mkdtemp()
        self.old_sys_path = list(sys.path)
        sys.path.insert(0, self.test_dir)
        self.shell = FakeShell()

    def teardown(self):
        for mod_name in self.created_temp_modules:
            sys.modules.pop(mod_name, None)
        shutil.rmtree(self.test_dir)
        sys.path = self.old_sys_path

    def get_module(self):
        module_name = "tmpmod_" + "".join(random.sample(self.filename_chars, 20))
        if module_name in sys.modules:
            del sys.modules[module_name]
        file_name = os.path.join(self.test_dir, module_name + ".py")
        return module_name, file_name

    def write_file(self, filename, content):
        """
        Write a file, and force a timestamp difference of at least one second

        Notes
        -----
        Python's .pyc files record the timestamp of their compilation
        with a time resolution of one second.

        Therefore, we need to force a timestamp difference between .py
        and .pyc, without having the .py file be timestamped in the
        future, and without changing the timestamp of the .pyc file
        (because that is stored in the file). The only reliable way
        to achieve this seems to be to sleep.

        Doesn't seem necessary on Darwin so we make this the exception.
        """
        if platform.system().lower() != "darwin":
            time.sleep(1.05)
        with open(filename, "w", encoding="utf-8") as f:
            f.write(squish_text(content))

    def new_module(self, code):
        mod_name, mod_fn = self.get_module()
        with open(mod_fn, "w", encoding="utf-8") as f:
            f.write(squish_text(code))
        self.created_temp_modules.add(mod_name)
        return mod_name, mod_fn


@pytest.fixture
def shell_fixture():
    helper = _ShellFixtureHelper()
    yield helper
    helper.teardown()


# ---------------------------------------------------------------------------
# AutoreloadHookSuite
# ---------------------------------------------------------------------------


@pytest.mark.parametrize("autoreload_when", [
    "before_import",
    "after_import",
    "after_first_pass",
])
def test_deduperreloader_basic(shell_fixture, autoreload_when):
    """Test autoreload works regardless of when magic_autoreload("2") is called."""
    src_v1 = """
        x = 9
        def foo(y):
            return y + 3
    """
    src_v2 = """
        x = 9
        def foo(y):
            return y + 5
        """
    if autoreload_when == "before_import":
        shell_fixture.shell.magic_autoreload("2")
    mod_name, mod_fn = shell_fixture.new_module(src_v1)
    shell_fixture.shell.run_code("import %s" % mod_name)
    if autoreload_when == "after_import":
        shell_fixture.shell.magic_autoreload("2")
    shell_fixture.shell.run_code("pass")
    if autoreload_when == "after_first_pass":
        shell_fixture.shell.magic_autoreload("2")
    shell_fixture.shell.run_code("pass")
    shell_fixture.write_file(mod_fn, src_v2)
    shell_fixture.shell.run_code("pass")
    assert mod_name in shell_fixture.shell.user_ns
    mod = sys.modules[mod_name]
    assert mod.foo(0) == 5


def test_super(shell_fixture):
    shell_fixture.shell.magic_autoreload("2")
    mod_name, mod_fn = shell_fixture.new_module(
        """
        class Foo:
            def foo(self):
                return 1
        class Bar(Foo):
            def bar(self):
                return super().foo() + 1
    """
    )
    shell_fixture.shell.run_code(f"from {mod_name} import Bar; bar = Bar()")
    shell_fixture.shell.run_code("pass")
    shell_fixture.write_file(
        mod_fn,
        """
        class Foo:
            def foo(self):
                return 1
        class Bar(Foo):
            def bar(self):
                return super().foo() + 2
        """,
    )
    shell_fixture.shell.run_code("result = bar.bar()")
    # NOTE : only works on CPython, requires read-only patching.
    if platform.python_implementation() == "CPython":
        assert shell_fixture.shell.user_ns["result"] == 3


def test_deduperreloader_need_to_default_back(shell_fixture):
    shell_fixture.shell.magic_autoreload("2")
    mod_name, mod_fn = shell_fixture.new_module(
        """
        x = 9
        def foo(y):
            return y + 3
    """
    )
    shell_fixture.shell.run_code("import %s" % mod_name)
    shell_fixture.shell.run_code("pass")
    shell_fixture.write_file(
        mod_fn,
        """
        x = 200
        def foo(y):
            return y + 5
        """,
    )
    shell_fixture.shell.run_code("pass")
    assert mod_name in shell_fixture.shell.user_ns
    mod = sys.modules[mod_name]
    assert mod.foo(0) == 5


def test_deduperreloader_failure(shell_fixture):
    shell_fixture.shell.magic_autoreload("2")
    mod_name, mod_fn = shell_fixture.new_module(
        """
        x = 9
        def foo(y):
            return y + 3
    """
    )
    shell_fixture.shell.run_code("import %s" % mod_name)
    shell_fixture.shell.run_code("pass")
    shell_fixture.write_file(
        mod_fn,
        """
        broken broken broken
        """,
    )
    shell_fixture.shell.run_code("pass")
    assert mod_name in shell_fixture.shell.user_ns
    mod = sys.modules[mod_name]
    assert mod.foo(0) == 3


def test_deduperreloader_imported_mod(shell_fixture):
    shell_fixture.shell.magic_autoreload("2")
    mod_name, mod_fn = shell_fixture.new_module(
        """
        from os import environ
        def foo(n):
            pass
    """
    )
    shell_fixture.shell.run_code("import %s" % mod_name)
    shell_fixture.shell.run_code("pass")
    shell_fixture.write_file(
        mod_fn,
        """
        from os import environ
        def foo():
            environ._data
            return 1
    """,
    )
    shell_fixture.shell.run_code("pass")
    assert mod_name in shell_fixture.shell.user_ns
    mod = sys.modules[mod_name]
    assert mod.foo() == 1


# ---------------------------------------------------------------------------
# AutoreloadClassMethodsDetectionSuite
# ---------------------------------------------------------------------------


def test_autoreload_method():
    code1 = squish_text(
        """
        class C(n):
            def foo():
                pass
    """
    )
    code2 = squish_text(
        """
        class C(n):
            def foo():
                x = 1
                return x
    """
    )
    deduperreloader = DeduperTestReloader()
    ast_1 = ast.parse(code1)
    ast_2 = ast.parse(code2)

    assert deduperreloader.detect_autoreload(ast_1, ast_2)
    assert deduperreloader._to_autoreload.defs_to_reload == []
    assert "C" in deduperreloader._to_autoreload.children
    assert ["foo"] == list(
        d[0][0] for d in deduperreloader._to_autoreload.children["C"].defs_to_reload
    )
    assert deduperreloader._to_autoreload.children["C"].children == {}


def test_autoreload_no_changes_class():
    code1 = squish_text(
        """
        class D(n):
            def foo():
                pass
        def foo():
            pass
    """
    )
    code2 = squish_text(
        """
        class D(n):
            def foo():
                pass
        def foo():
            pass
    """
    )
    deduperreloader = DeduperTestReloader()
    ast_1 = ast.parse(code1)
    ast_2 = ast.parse(code2)

    assert deduperreloader.detect_autoreload(ast_1, ast_2)
    assert deduperreloader._to_autoreload.defs_to_reload == []
    assert deduperreloader._to_autoreload.children == {}


def test_autoreload_add_function_class():
    code1 = squish_text(
        """
        class C:
            def foo():
                pass
        def foo():
            pass
    """
    )
    code2 = squish_text(
        """
        class C:
            def foo():
                x = 1
                return x
            def bar():
                return -1
        def foo():
            pass
    """
    )
    deduperreloader = DeduperTestReloader()
    ast_1 = ast.parse(code1)
    ast_2 = ast.parse(code2)

    assert deduperreloader.detect_autoreload(ast_1, ast_2)
    assert deduperreloader._to_autoreload.defs_to_reload == []
    assert "C" in deduperreloader._to_autoreload.children
    assert ["foo", "bar"] == list(
        d[0][0] for d in deduperreloader._to_autoreload.children["C"].defs_to_reload
    )
    assert len(deduperreloader._to_autoreload.children["C"].children) == 0


def test_autoreload_remove_method():
    code1 = squish_text(
        """
        class C:
            def foo():
                pass
        def foo():
            pass
    """
    )
    code2 = squish_text(
        """
        class C:
            pass
        def foo():
            pass
    """
    )
    deduperreloader = DeduperTestReloader()
    ast_1 = ast.parse(code1)
    ast_2 = ast.parse(code2)

    assert deduperreloader.detect_autoreload(ast_1, ast_2)


def test_autoreload_remove_class():
    code1 = squish_text(
        """
        class C:
            def foo():
                pass
        def foo():
            pass
    """
    )
    code2 = squish_text(
        """
        def foo():
            pass
    """
    )
    deduperreloader = DeduperTestReloader()
    ast_1 = ast.parse(code1)
    ast_2 = ast.parse(code2)

    assert deduperreloader.detect_autoreload(ast_1, ast_2)


def test_autoreload_add_class_in_class():
    code1 = squish_text(
        """
        class C:
            pass
        def foo():
            pass
    """
    )
    code2 = squish_text(
        """
        class C:
            class D:
                pass
        def foo():
            pass
    """
    )
    deduperreloader = DeduperTestReloader()
    ast_1 = ast.parse(code1)
    ast_2 = ast.parse(code2)

    assert deduperreloader.detect_autoreload(ast_1, ast_2)


def test_autoreload_add_method_in_class_in_class():
    code1 = squish_text(
        """
        class C:
            class D:
                x = 1
        def foo():
            pass
    """
    )
    code2 = squish_text(
        """
        class C:
            class D:
                x = 1
                def foo():
                    return 1
        def foo():
            pass
    """
    )
    deduperreloader = DeduperTestReloader()
    ast_1 = ast.parse(code1)
    ast_2 = ast.parse(code2)

    assert deduperreloader.detect_autoreload(ast_1, ast_2)
    assert deduperreloader._to_autoreload.defs_to_reload == []
    assert list(deduperreloader._to_autoreload.children.keys()) == ["C"]
    assert deduperreloader._to_autoreload.children["C"].defs_to_reload == []
    assert list(deduperreloader._to_autoreload.children["C"].children.keys()) == [
        "D"
    ]
    assert list(
        d[0][0]
        for d in deduperreloader._to_autoreload.children["C"]
        .children["D"]
        .defs_to_reload
    ) == ["foo"]
    assert deduperreloader._to_autoreload.children["C"].children["D"].children == {}


def test_autoreload_add_var_and_method_in_class_in_class():
    code1 = squish_text(
        """
        class C:
            class D:
                pass
        def foo():
            pass
    """
    )
    code2 = squish_text(
        """
        class C:
            class D:
                x = 1
                def foo():
                    return 1
        def foo():
            pass
    """
    )
    deduperreloader = DeduperTestReloader()
    ast_1 = ast.parse(code1)
    ast_2 = ast.parse(code2)

    assert deduperreloader.detect_autoreload(ast_1, ast_2)


def test_autoreload_add_method_in_class_in_class_pass():
    code1 = squish_text(
        """
        class C:
            class D:
                pass
        def foo():
            pass
    """
    )
    code2 = squish_text(
        """
        class C:
            class D:
                def foo():
                    return 1
        def foo():
            pass
    """
    )
    deduperreloader = DeduperTestReloader()
    ast_1 = ast.parse(code1)
    ast_2 = ast.parse(code2)

    assert deduperreloader.detect_autoreload(ast_1, ast_2)
    assert deduperreloader._to_autoreload.defs_to_reload == []
    assert list(deduperreloader._to_autoreload.children.keys()) == ["C"]
    assert deduperreloader._to_autoreload.children["C"].defs_to_reload == []
    assert list(deduperreloader._to_autoreload.children["C"].children.keys()) == [
        "D"
    ]
    assert list(
        d[0][0]
        for d in deduperreloader._to_autoreload.children["C"]
        .children["D"]
        .defs_to_reload
    ) == ["foo"]
    assert deduperreloader._to_autoreload.children["C"].children["D"].children == {}


def test_autoreload_add_method_in_class_in_class_more():
    code1 = squish_text(
        """
        class C:
            class D:
                pass
        def foo():
            pass
    """
    )
    code2 = squish_text(
        """
        class C:
            def bar():
                return -1
            class D:
                def foo():
                    return 1
        def foo():
            pass
    """
    )
    deduperreloader = DeduperTestReloader()
    ast_1 = ast.parse(code1)
    ast_2 = ast.parse(code2)

    assert deduperreloader.detect_autoreload(ast_1, ast_2)
    assert deduperreloader._to_autoreload.defs_to_reload == []
    assert list(deduperreloader._to_autoreload.children.keys()) == ["C"]
    assert list(
        d[0][0] for d in deduperreloader._to_autoreload.children["C"].defs_to_reload
    ) == ["bar"]
    assert list(deduperreloader._to_autoreload.children["C"].children.keys()) == [
        "D"
    ]
    assert list(
        d[0][0]
        for d in deduperreloader._to_autoreload.children["C"]
        .children["D"]
        .defs_to_reload
    ) == ["foo"]
    assert deduperreloader._to_autoreload.children["C"].children["D"].children == {}


def test_autoreload_add_method_in_class_in_class_with_members():
    code1 = squish_text(
        """
        class C:
            class D:
                pass
        def foo():
            pass
    """
    )
    code2 = squish_text(
        """
        class C:
            def bar():
                return -1
            class D:
                s = 2
                def foo():
                    return 1
        def foo():
            pass
    """
    )
    deduperreloader = DeduperTestReloader()
    ast_1 = ast.parse(code1)
    ast_2 = ast.parse(code2)

    assert deduperreloader.detect_autoreload(ast_1, ast_2)


# ---------------------------------------------------------------------------
# AutoreloadReliabilitySuite
# ---------------------------------------------------------------------------


def test_autoreload_class_basic(shell_fixture):
    shell_fixture.shell.magic_autoreload("2")
    mod_name, mod_fn = shell_fixture.new_module(
        """
        x = 9
        class C:
            def foo():
                return 1
    """
    )
    shell_fixture.shell.run_code("import %s" % mod_name)
    shell_fixture.shell.run_code("pass")
    shell_fixture.write_file(
        mod_fn,
        """
        x = 9
        class C:
            def foo():
                return 1
        """,
    )
    shell_fixture.shell.run_code("pass")
    assert mod_name in shell_fixture.shell.user_ns
    mod = sys.modules[mod_name]
    assert mod.C.foo() == 1


def test_remove_overridden_method(shell_fixture):
    shell_fixture.shell.magic_autoreload("2")
    mod_name, mod_fn = shell_fixture.new_module(
        """
        class A:
            def foo(self):
                return 1
        class B(A):
            def foo(self):
                return 42
    """
    )
    shell_fixture.shell.run_code(f"from {mod_name} import B; b = B()")
    shell_fixture.shell.run_code("assert b.foo() == 42")
    shell_fixture.write_file(
        mod_fn,
        """
        class A:
            def foo(self):
                return 1
        class B(A):
            pass
        """,
    )
    shell_fixture.shell.run_code("assert b.foo() == 1")


def test_autoreload_class_use_outside_func(shell_fixture):
    shell_fixture.shell.magic_autoreload("2")
    mod_name, mod_fn = shell_fixture.new_module(
        """
        x = 9
        class C:
            def foo():
                return 1
    """
    )
    shell_fixture.shell.run_code("import %s" % mod_name)
    shell_fixture.shell.run_code("pass")
    shell_fixture.write_file(
        mod_fn,
        """
        x = 9
        class C:
            def foo():
                return 1+x
        """,
    )
    shell_fixture.shell.run_code("pass")
    assert mod_name in shell_fixture.shell.user_ns
    mod = sys.modules[mod_name]
    assert mod.C.foo() == 10


def test_autoreload_class_use_class_member(shell_fixture):
    shell_fixture.shell.magic_autoreload("2")
    mod_name, mod_fn = shell_fixture.new_module(
        """
        x = 9
        class C:
            def foo():
                return 1
    """
    )
    shell_fixture.shell.run_code("import %s" % mod_name)
    shell_fixture.shell.run_code("pass")
    shell_fixture.write_file(
        mod_fn,
        """
        class C:
            x = 9
            def foo():
                return 1+C.x
        """,
    )
    shell_fixture.shell.run_code("pass")
    assert mod_name in shell_fixture.shell.user_ns
    mod = sys.modules[mod_name]
    assert mod.C.foo() == 10


def test_autoreload_class_pass(shell_fixture):
    shell_fixture.shell.magic_autoreload("2")
    mod_name, mod_fn = shell_fixture.new_module(
        """
        x = 9
        class C:
            pass
    """
    )
    shell_fixture.shell.run_code("import %s" % mod_name)
    shell_fixture.shell.run_code("pass")
    shell_fixture.write_file(
        mod_fn,
        """
        class C:
            x = 9
            def foo():
                return 1+C.x
        """,
    )
    shell_fixture.shell.run_code("pass")
    assert mod_name in shell_fixture.shell.user_ns
    mod = sys.modules[mod_name]
    assert mod.C.foo() == 10


def test_autoreload_class_ellipsis(shell_fixture):
    shell_fixture.shell.magic_autoreload("2")
    mod_name, mod_fn = shell_fixture.new_module(
        """
        x = 9
        class C:
            ...
    """
    )
    shell_fixture.shell.run_code("import %s" % mod_name)
    shell_fixture.shell.run_code("pass")
    shell_fixture.write_file(
        mod_fn,
        """
        class C:
            x = 9
            def foo():
                return 1+C.x
        """,
    )
    shell_fixture.shell.run_code("pass")
    assert mod_name in shell_fixture.shell.user_ns
    mod = sys.modules[mod_name]
    assert mod.C.foo() == 10


def test_autoreload_class_default_autoreload(shell_fixture):
    shell_fixture.shell.magic_autoreload("2")
    mod_name, mod_fn = shell_fixture.new_module(
        """
        class C:
            x = 9
            def foo():
                return 1+C.x
    """
    )
    shell_fixture.shell.run_code("import %s" % mod_name)
    shell_fixture.shell.run_code("pass")
    shell_fixture.write_file(
        mod_fn,
        """
        class C:
            x = 20
            def foo():
                return 1+C.x
        """,
    )
    shell_fixture.shell.run_code("pass")
    assert mod_name in shell_fixture.shell.user_ns
    mod = sys.modules[mod_name]
    assert mod.C.foo() == 21


def test_autoreload_class_nested(shell_fixture):
    shell_fixture.shell.magic_autoreload("2")
    mod_name, mod_fn = shell_fixture.new_module(
        """
        class C:
            x = 9
            class D:
                def foo():
                    pass
    """
    )
    shell_fixture.shell.run_code("import %s" % mod_name)
    shell_fixture.shell.run_code("pass")
    shell_fixture.write_file(
        mod_fn,
        """
        class C:
            x = 9
            class D:
                def foo():
                    return 10
        """,
    )
    shell_fixture.shell.run_code("pass")
    assert mod_name in shell_fixture.shell.user_ns
    mod = sys.modules[mod_name]
    assert mod.C.D.foo() == 10


def test_autoreload_class_nested2(shell_fixture):
    shell_fixture.shell.magic_autoreload("2")
    mod_name, mod_fn = shell_fixture.new_module(
        """
        class C:
            x = 9
            def c():
                return 1
            class D:
                def foo():
                    pass
    """
    )
    shell_fixture.shell.run_code("import %s" % mod_name)
    shell_fixture.shell.run_code("pass")
    shell_fixture.write_file(
        mod_fn,
        """
        class C:
            x = 9
            def c():
                return 1
            class D:
                def foo():
                    return 10
        """,
    )
    shell_fixture.shell.run_code("pass")
    assert mod_name in shell_fixture.shell.user_ns
    mod = sys.modules[mod_name]
    assert mod.C.D.foo() == 10
    assert mod.C.c() == 1


def test_autoreload_class_nested3(shell_fixture):
    shell_fixture.shell.magic_autoreload("2")
    mod_name, mod_fn = shell_fixture.new_module(
        """
        class C:
            x = 9
            def c():
                return 1
            class D:
                def foo():
                    pass
    """
    )
    shell_fixture.shell.run_code("import %s" % mod_name)
    shell_fixture.shell.run_code("pass")
    shell_fixture.write_file(
        mod_fn,
        """
        class C:
            x = 9
            def c():
                return 13
            class D:
                def foo():
                    return 10
        """,
    )
    shell_fixture.shell.run_code("pass")
    assert mod_name in shell_fixture.shell.user_ns
    mod = sys.modules[mod_name]
    assert mod.C.D.foo() == 10
    assert mod.C.c() == 13


def test_autoreload_new_class_added(shell_fixture):
    shell_fixture.shell.magic_autoreload("2")
    mod_name, mod_fn = shell_fixture.new_module(
        """
        class C:
            x = 9
            def c():
                return 1
    """
    )
    shell_fixture.shell.run_code("import %s" % mod_name)
    shell_fixture.shell.run_code("pass")
    shell_fixture.write_file(
        mod_fn,
        """
        class C:
            x = 9
            def c():
                return 13
        class D:
            def c():
                return 1
        """,
    )
    shell_fixture.shell.run_code("pass")
    assert mod_name in shell_fixture.shell.user_ns
    mod = sys.modules[mod_name]
    assert mod.C.c() == 13
    assert mod.D.c() == 1


def test_autoreload_class_nested_default(shell_fixture):
    shell_fixture.shell.magic_autoreload("2")
    mod_name, mod_fn = shell_fixture.new_module(
        """
        class C:
            pass
    """
    )
    shell_fixture.shell.run_code("import %s" % mod_name)
    shell_fixture.shell.run_code("pass")
    shell_fixture.write_file(
        mod_fn,
        """
        class C:
            x = 9
            def c():
                return 13
            class D:
                def foo():
                    return 10
        """,
    )
    shell_fixture.shell.run_code("pass")
    assert mod_name in shell_fixture.shell.user_ns
    mod = sys.modules[mod_name]
    assert mod.C.D.foo() == 10
    assert mod.C.c() == 13


def test_autoreload_class_nested_using_members(shell_fixture):
    shell_fixture.shell.magic_autoreload("2")
    mod_name, mod_fn = shell_fixture.new_module(
        """
        class C:
            x = 9
            def c():
                return 13
            class D:
                def foo():
                    return 10
        """
    )
    shell_fixture.shell.run_code("import %s" % mod_name)
    shell_fixture.shell.run_code("pass")
    shell_fixture.write_file(
        mod_fn,
        """
        class C:
            x = 9
            def c():
                return 13
            class D:
                def foo():
                    return 10 + C.x
        """,
    )
    shell_fixture.shell.run_code("pass")
    assert mod_name in shell_fixture.shell.user_ns
    mod = sys.modules[mod_name]
    assert mod.C.D.foo() == 19
    assert mod.C.c() == 13


def test_autoreload_class_nested_using_members_ellipsis(shell_fixture):
    shell_fixture.shell.magic_autoreload("2")
    mod_name, mod_fn = shell_fixture.new_module(
        """
        class C:
            x = 9
            def c():
                return 13
            class D:
                def foo():
                    ...
        """
    )
    shell_fixture.shell.run_code("import %s" % mod_name)
    shell_fixture.shell.run_code("pass")
    shell_fixture.write_file(
        mod_fn,
        """
        class C:
            x = 9
            def c():
                return 13
            class D:
                def foo():
                    return 10 + C.x
        """,
    )
    shell_fixture.shell.run_code("pass")
    assert mod_name in shell_fixture.shell.user_ns
    mod = sys.modules[mod_name]
    assert mod.C.D.foo() == 19
    assert mod.C.c() == 13


def test_method_decorators_no_changes(shell_fixture):
    shell_fixture.shell.magic_autoreload("2")
    mod_name, mod_file = shell_fixture.new_module(
        """
        class Foo:
            @classmethod
            def bar(cls):
                return 0

            @classmethod
            def foo(cls):
                return 42 + cls.bar()

        foo = Foo.foo
        """
    )
    shell_fixture.shell.run_code("import %s" % mod_name)
    shell_fixture.shell.run_code("pass")
    assert mod_name in shell_fixture.shell.user_ns
    mod = sys.modules[mod_name]
    assert mod.foo() == 42

    shell_fixture.shell.run_code(f"assert {mod_name}.foo() == 42")
    shell_fixture.write_file(
        mod_file,
        """
        class Foo:
            @classmethod
            def bar(cls):
                return 0

            @classmethod
            def foo(cls):
                return 42 + cls.bar()

        foo = Foo.foo
        """,
    )
    shell_fixture.shell.run_code(f"assert {mod_name}.foo() == 42")


def test_method_decorators_no_changes1(shell_fixture):
    shell_fixture.shell.magic_autoreload("2")
    mod_name, mod_file = shell_fixture.new_module(
        """
        class Foo:
            @classmethod
            def bar(cls):
                return 0

            @classmethod
            def foo(cls):
                return 42 + cls.bar()

        foo = Foo.foo
        """
    )

    shell_fixture.shell.run_code(f"from {mod_name} import foo")
    shell_fixture.shell.run_code("assert foo() == 42")
    shell_fixture.write_file(
        mod_file,
        """
        class Foo:
            @classmethod
            def bar(cls):
                return 0

            @classmethod
            def foo(cls):
                return 42 + cls.bar()

        foo = Foo.foo
        """,
    )
    shell_fixture.shell.run_code("assert foo() == 42")


def test_method_classmethod_one_change(shell_fixture):
    shell_fixture.shell.magic_autoreload("2")
    mod_name, mod_file = shell_fixture.new_module(
        """
        class Foo:
            @classmethod
            def bar(cls):
                return 0

            @classmethod
            def func(cls):
                return 42 + cls.bar()

        func = Foo.func
        """
    )
    shell_fixture.shell.run_code("import %s" % mod_name)
    shell_fixture.shell.run_code(f"assert {mod_name}.func() == 42")
    shell_fixture.write_file(
        mod_file,
        """
        class Foo:
            @classmethod
            def bar(cls):
                return 1

            @classmethod
            def func(cls):
                return 42 + cls.bar()

        func = Foo.func
        """,
    )
    shell_fixture.shell.run_code(f"assert {mod_name}.func() == 43")


def test_method_staticmethod_one_change(shell_fixture):
    shell_fixture.shell.magic_autoreload("2")
    mod_name, mod_file = shell_fixture.new_module(
        """
        class Foo:
            @staticmethod
            def bar():
                return 0

            @staticmethod
            def func():
                return 42 + Foo.bar()

        func = Foo.func
        """
    )
    shell_fixture.shell.run_code(f"from {mod_name} import func")
    shell_fixture.shell.run_code("assert func() == 42")
    shell_fixture.write_file(
        mod_file,
        """
        class Foo:
            @staticmethod
            def bar():
                return 1

            @staticmethod
            def func():
                return 42 + Foo.bar()

        func = Foo.func
        """,
    )
    shell_fixture.shell.run_code("assert func() == 43")


def test_autoreload_class_default_args(shell_fixture):
    shell_fixture.shell.magic_autoreload("2")
    mod_name, mod_fn = shell_fixture.new_module(
        """
        x = 42
        class Foo:
            def foo(self, y): return y
    """
    )
    shell_fixture.shell.run_code("import %s" % mod_name)
    shell_fixture.shell.run_code("pass")
    shell_fixture.write_file(
        mod_fn,
        """
        x = 42
        class Foo:
            def foo(self, y=x): return y
        """,
    )
    shell_fixture.shell.run_code("pass")
    assert mod_name in shell_fixture.shell.user_ns
    mod = sys.modules[mod_name]
    obj = mod.Foo()
    assert obj.foo(2) == 2
    assert obj.foo() == 42


def test_autoreload_class_change_default_args(shell_fixture):
    shell_fixture.shell.magic_autoreload("2")
    mod_name, mod_fn = shell_fixture.new_module(
        """
        x = 42
        class Foo:
            def foo(y): return y
    """
    )
    shell_fixture.shell.run_code("import %s" % mod_name)
    shell_fixture.shell.run_code("pass")
    shell_fixture.write_file(
        mod_fn,
        """
        x = 44
        class Foo:
            def foo(y=x): return y
        """,
    )
    shell_fixture.shell.run_code("pass")
    assert mod_name in shell_fixture.shell.user_ns
    mod = sys.modules[mod_name]
    assert mod.Foo.foo() == 44


def test_autoreload_class_new_class(shell_fixture):
    shell_fixture.shell.magic_autoreload("2")
    mod_name, mod_fn = shell_fixture.new_module(
        """
        x = 42
        class Foo:
            def foo(y=x): return y
    """
    )
    shell_fixture.shell.run_code("import %s" % mod_name)
    shell_fixture.shell.run_code("pass")
    prev_foo = shell_fixture.shell.user_ns[mod_name].Foo.foo
    shell_fixture.write_file(
        mod_fn,
        """
        x = 42
        class Foo:
            def foo(y=x): return y
        class C:
            def foo(): return 200
        """,
    )
    shell_fixture.shell.run_code("pass")
    assert mod_name in shell_fixture.shell.user_ns
    assert prev_foo is shell_fixture.shell.user_ns[mod_name].Foo.foo
    mod = sys.modules[mod_name]
    assert mod.Foo.foo() == 42
    assert mod.C.foo() == 200


def test_autoreload_overloaded_vars(shell_fixture):
    shell_fixture.shell.magic_autoreload("2")
    mod_name, mod_fn = shell_fixture.new_module(
        """
        x = 42
        class Foo:
            pass
    """
    )
    shell_fixture.shell.run_code("import %s" % mod_name)
    shell_fixture.shell.run_code("pass")
    mod = sys.modules[mod_name]
    shell_fixture.write_file(
        mod_fn,
        """
        x = 44
        class Foo:
            def foo():
                x = 2
                return x
        """,
    )
    shell_fixture.shell.run_code("pass")
    assert mod_name in shell_fixture.shell.user_ns
    assert mod.Foo.foo() == 2


def test_autoreload_overloaded_vars2(shell_fixture):
    shell_fixture.shell.magic_autoreload("2")
    mod_name, mod_fn = shell_fixture.new_module(
        """
        x = 42
        def foo():
            return x
    """
    )
    shell_fixture.shell.run_code("import %s" % mod_name)
    shell_fixture.shell.run_code("pass")
    mod = sys.modules[mod_name]
    shell_fixture.write_file(
        mod_fn,
        """
        x = 44
        def foo():
            x = 2
            return x
        """,
    )
    shell_fixture.shell.run_code("pass")
    assert mod_name in shell_fixture.shell.user_ns
    assert mod.foo() == 2


# ---------------------------------------------------------------------------
# DecoratorPatchingSuite
# ---------------------------------------------------------------------------


def test_modify_property(shell_fixture):
    shell_fixture.shell.magic_autoreload("2")
    mod_name, mod_file = shell_fixture.new_module(
        """
        class Foo:
            @property
            def foo(self):
                return 42
        """
    )
    shell_fixture.shell.run_code(f"from {mod_name} import Foo")
    shell_fixture.shell.run_code("foo = Foo()")
    shell_fixture.shell.run_code("assert foo.foo == 42")
    shell_fixture.write_file(
        mod_file,
        """
        class Foo:
            @property
            def foo(self):
                return 43
        """,
    )
    shell_fixture.shell.run_code("pass")
    shell_fixture.shell.run_code("assert foo.foo == 43")


def test_method_decorator(shell_fixture):
    shell_fixture.shell.magic_autoreload("2")
    mod_name, mod_file = shell_fixture.new_module(
        """
        def incremented(f):
            return lambda *args: f(*args) + 1

        class Foo:
            @classmethod
            @incremented
            def foo(cls):
                return 42

        foo = Foo.foo
        """
    )
    shell_fixture.shell.run_code(f"from {mod_name} import foo")
    shell_fixture.shell.run_code("assert foo() == 43")
    shell_fixture.write_file(
        mod_file,
        """
        def incremented(f):
            return lambda *args: f(*args) + 1

        class Foo:
            @classmethod
            def foo(cls):
                return 42

        foo = Foo.foo
        """,
    )
    shell_fixture.shell.run_code("assert foo() == 42")


def test_method_modified_decorator(shell_fixture):
    shell_fixture.shell.magic_autoreload("2")
    mod_name, mod_file = shell_fixture.new_module(
        """
        def incremented(f):
            return lambda *args: f(*args) + 1

        class Foo:
            @classmethod
            @incremented
            def foo(cls):
                return 42

        foo = Foo.foo
        """
    )
    shell_fixture.shell.run_code(f"from {mod_name} import foo")
    shell_fixture.shell.run_code("assert foo() == 43")
    shell_fixture.write_file(
        mod_file,
        """
        def incremented(f):
            return lambda *args: f(*args) + 0

        class Foo:
            @classmethod
            @incremented
            def foo(cls):
                return 42

        foo = Foo.foo
        """,
    )
    shell_fixture.shell.run_code("assert foo() == 42")


def test_function_decorators(shell_fixture):
    shell_fixture.shell.magic_autoreload("2")
    mod_name, mod_file = shell_fixture.new_module(
        """
        def incremented(f):
            return lambda *args: f(*args) + 1

        @incremented
        def foo():
            return 42
        """
    )
    shell_fixture.shell.run_code("import %s" % mod_name)
    shell_fixture.shell.run_code("pass")
    mod = sys.modules[mod_name]
    assert mod.foo() == 43
    shell_fixture.write_file(
        mod_file,
        """
        def incremented(f):
            return lambda *args: f(*args) + 1

        def foo():
            return 42
        """,
    )
    shell_fixture.shell.run_code("pass")
    assert mod.foo() == 42
    shell_fixture.write_file(
        mod_file,
        """
        def incremented(v):
            def deco(f):
                return lambda *args: f(*args) + v
            return deco

        @incremented(2)
        def foo():
            return 43
        """,
    )
    shell_fixture.shell.run_code("pass")
    assert mod.foo() == 45


def test_method_decorators_again(shell_fixture):
    shell_fixture.shell.magic_autoreload("2")
    mod_name, mod_file = shell_fixture.new_module(
        """
        class Foo:
            @classmethod
            def bar(cls):
                return 0

            @classmethod
            def foo(cls):
                return 42 + cls.bar()

        foo = Foo.foo
        """
    )
    shell_fixture.shell.run_code("import %s" % mod_name)
    shell_fixture.shell.run_code("pass")
    mod = sys.modules[mod_name]
    assert mod.foo() == 42
    shell_fixture.write_file(
        mod_file,
        """
        class Foo:
            @classmethod
            def bar(cls):
                return 1

            @classmethod
            def foo(cls):
                return 42 + cls.bar()

        foo = Foo.foo
        """,
    )
    shell_fixture.shell.run_code("pass")
    assert mod.Foo.foo() == 43
    assert mod.foo() == 43


# This test verifies that the correct globals are used when patching a function
# decorated by a function from another module. For this example, we should always
# use <mod_name>.__dict__ as the global environment when patching foo, which comes
# from <mod_name>, and never <other_mod_name>.__dict__, which is what we would get
# if we use foo.__globals__ after it has been decorated.
def test_function_decorator_from_other_module(shell_fixture):
    shell_fixture.shell.magic_autoreload("2")
    other_mod_name, _ = shell_fixture.new_module(
        """
        import functools

        def incremented(f):
            @functools.wraps(f)
            def wrapper(*args, **kwargs):
                return f(*args, **kwargs) + 1
            return wrapper
        """
    )
    mod_name, mod_file = shell_fixture.new_module(
        f"""
        from {other_mod_name} import incremented as deco
        @deco
        def foo():
            return 42
        """
    )
    shell_fixture.shell.run_code(f"from {mod_name} import foo")
    shell_fixture.shell.run_code("assert foo() == 43")
    shell_fixture.write_file(
        mod_file,
        f"""
        from {other_mod_name} import incremented as deco
        @deco
        def foo():
            return 43
        """,
    )
    shell_fixture.shell.run_code("assert foo() == 44")


def test_decorators_with_freevars(shell_fixture):
    shell_fixture.shell.magic_autoreload("2")
    mod_name, mod_file = shell_fixture.new_module(
        """
        def decorate(func):
            free_var_x = "x"
            free_var_y = "y"
            def wrapper(*args, **kwargs):
                return func(*args, **kwargs), free_var_x, free_var_y
            return wrapper
        @decorate
        def f():
            return "v1"
        """
    )
    shell_fixture.shell.run_code(f"from {mod_name} import f")
    mod = sys.modules[mod_name]
    assert mod.f() == ("v1", "x", "y")
    shell_fixture.write_file(
        mod_file,
        """
        def decorate(func):
            free_var_ax = "ax"
            free_var_by = "by"
            def wrapper(*args, **kwargs):
                return func(*args, **kwargs), free_var_ax, free_var_by
            return wrapper
        @decorate
        def f():
            return "v2"
        """,
    )
    shell_fixture.shell.run_code("pass")
    val = mod.f()
    assert val == ("v2", "ax", "by"), val


# ---------------------------------------------------------------------------
# TestAutoreloadEnum
# ---------------------------------------------------------------------------


def test_reload_enums(shell_fixture):
    shell_fixture.shell.magic_autoreload("2")
    mod_name, mod_fn = shell_fixture.new_module(
        """
        from enum import Enum
        class MyEnum(Enum):
            A = 'A'
            B = 'B'
        """,
    )
    shell_fixture.shell.run_code("import %s" % mod_name)
    shell_fixture.shell.run_code("pass")
    mod = sys.modules[mod_name]
    shell_fixture.write_file(
        mod_fn,
        """
        from enum import Enum
        class MyEnum(Enum):
            A = 'A'
            B = 'B'
            C = 'C'
        """,
    )
    shell_fixture.shell.run_code("pass")
    assert mod.MyEnum.C.value == "C"
