diff --git a/src/patchy/api.py b/src/patchy/api.py index d014aed..d87aca6 100644 --- a/src/patchy/api.py +++ b/src/patchy/api.py @@ -299,11 +299,18 @@ def _process_function() -> ast.Module: # Compile and retrieve the new Code object localz: dict[str, Any] = {} + + globalz = func.__globals__.copy() + for cell in func.__closure__ or (): + cell_contents = cell.cell_contents + if isinstance(cell_contents, type): + globalz.update(cell_contents.__dict__) + new_code = cast(CodeType, _compile(new_source)) exec( new_code, - dict(func.__globals__), + globalz, localz, ) new_func = localz["__patchy_freevars__"]() diff --git a/tests/test_patch.py b/tests/test_patch.py index 785b030..e5ddc11 100644 --- a/tests/test_patch.py +++ b/tests/test_patch.py @@ -795,6 +795,49 @@ def bark() -> str: assert Doge.bark() == "Wowowow" +def test_patch_closure_use_function(): + DEFAULT_MEDIUM = "Chalk" + + def paint(medium: str = DEFAULT_MEDIUM) -> str: + return medium + + assert paint() == "Chalk" + + patchy.patch( + paint, + """\ + @@ -1,2 +1,2 @@ + def paint(medium: str = DEFAULT_MEDIUM) -> str: + - return medium + + return "Cheese" + """, + ) + + assert paint() == "Cheese" + + +def test_patch_closure_use_instancemethod(): + class Artist: + DEFAULT_MEDIUM = "Chalk" + + def paint(self, medium: str = DEFAULT_MEDIUM) -> str: + return medium + + assert Artist().paint() == "Chalk" + + patchy.patch( + Artist.paint, + """\ + @@ -1,2 +1,2 @@ + def paint(self, medium: str = DEFAULT_MEDIUM) -> str: + - return medium + + return "Cheese" + """, + ) + + assert Artist().paint() == "Cheese" + + def test_patch_future_python(tmp_path): (tmp_path / "future_user.py").write_text( dedent(