refactor: rename write_patched_source to persist_patched_source and update related usages

Author: fanck0605

src/awepatch/function.py

@@ -8,7 +8,7 @@
 from copy import deepcopy
 from functools import partial
 from types import CodeType, TracebackType
-from typing import TYPE_CHECKING, Any, Self
+from typing import Any, Self
 
 from awepatch.utils import (
     AWEPATCH_DEBUG,
@@ -22,9 +22,11 @@
     compile_idents,
     find_matched_node,
     load_stmts,
-    write_patched_source,
+    persist_patched_source,
 )
 
+TYPE_CHECKING = False
+
 if TYPE_CHECKING:
     from collections.abc import Callable, Sequence
 
@@ -88,7 +90,7 @@ def load_function_code(
     source = ast.unparse(func)
 
     if AWEPATCH_DEBUG:
-        file_path, source = write_patched_source(
+        file_path, source = persist_patched_source(
             source,
             name=func.name,
             type="function",
@@ -107,13 +109,13 @@ def load_function_code(
 
 
 def _get_function_def(
-    func: CodeType, source: list[str]
+    func: CodeType, slines: list[str]
 ) -> ast.FunctionDef | ast.AsyncFunctionDef:
     """Get the AST function definition from a code object.
 
     Args:
         func: The code object
-        source: The source code lines of the function
+        slines: The source code lines of the function
 
     Returns:
         The AST function definition
@@ -123,7 +125,7 @@ def _get_function_def(
 
     """
 
-    for node in ast.walk(ast.parse("".join(source))):
+    for node in ast.walk(ast.parse("".join(slines))):
         if not isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
             continue
         if node.name != func.co_name:
@@ -161,8 +163,8 @@ class _SingleFunctionPatcher:
     def __init__(self, func: Callable[..., Any]) -> None:
         self._func = func
         self._orig_code = func.__code__
-        self._source, _ = inspect.findsource(func)
-        self._func_def = _get_function_def(func.__code__, self._source)
+        self._slines, _ = inspect.findsource(func)
+        self._func_def = _get_function_def(func.__code__, self._slines)
         self._func_def.decorator_list.clear()
         self._patches: list[CompiledPatch] = []
 
@@ -183,11 +185,9 @@ def add_patch(
     def apply(self) -> Callable[..., Any]:
         """Apply the patches to the function."""
         func_def = deepcopy(self._func_def)
-        compiled: CompiledPatches = defaultdict(
-            lambda: defaultdict(lambda: defaultdict(list))
-        )
+        compiled: CompiledPatches = defaultdict(lambda: defaultdict(dict))
         for patch in self._patches:
-            target = find_matched_node(func_def, self._source, patch.target)
+            target = find_matched_node(func_def, self._slines, patch.target)
             if target is None:
                 raise ValueError(f"Patch target {patch.target} not found")
             append_patch(

src/awepatch/module.py

@@ -19,7 +19,7 @@
     compile_idents,
     find_matched_node,
     load_stmts,
-    write_patched_source,
+    persist_patched_source,
 )
 
 TYPE_CHECKING = False
@@ -51,9 +51,7 @@ def get_data(self, path: str) -> bytes:
             source = f.read()
         tree = ast.parse(source, filename=self._origin)
         slines = source.splitlines(keepends=True)
-        compiled: CompiledPatches = defaultdict(
-            lambda: defaultdict(lambda: defaultdict(list))
-        )
+        compiled: CompiledPatches = defaultdict(lambda: defaultdict(dict))
         for patch in self._patches:
             target = find_matched_node(tree, slines, patch.target)
             if target is None:
@@ -69,7 +67,7 @@ def get_data(self, path: str) -> bytes:
         apply_compiled_patches(compiled)
         source = ast.unparse(tree)
         if AWEPATCH_DEBUG:
-            self._path, source = write_patched_source(
+            self._path, source = persist_patched_source(
                 source,
                 self._fullname.rsplit(".", 1)[-1],
                 "module",

src/awepatch/utils.py

@@ -5,18 +5,15 @@
 import re
 import threading
 from abc import ABC, abstractmethod
-from binascii import crc32
 from collections import defaultdict
+from collections.abc import Sequence
 from dataclasses import KW_ONLY, dataclass
 from typing import Any, Literal, TypeAlias, cast
 
-from filelock import FileLock
-from platformdirs import user_cache_dir
-
 TYPE_CHECKING = False
 
 if TYPE_CHECKING:
-    from collections.abc import Mapping, Sequence
+    from collections.abc import Mapping
     from types import TracebackType
 
 
@@ -101,7 +98,7 @@ def __str__(self) -> str:
 
 
 CompiledPatches: TypeAlias = defaultdict[
-    Location, defaultdict[int, defaultdict[Mode, list[ast.stmt]]]
+    Location, defaultdict[int, dict[Mode, Sequence[ast.stmt]]]
 ]
 
 
@@ -112,19 +109,16 @@ def append_patch(
     mode: Mode,
 ) -> None:
     patches = complied[target[0]][target[1]]
-    # Check for conflicting patches,
-    if patches and mode == "replace" and "replace" in patches:
-        raise ValueError(f"Cannot have multiple 'replace' patches on target {target!r}")
     if mode == "replace":
         if "replace" in patches:
             raise ValueError(
                 f"Multiple 'replace' patches on the same target {target!r}"
             )
-        patches[mode].extend(stmts)
+        patches[mode] = stmts
     elif mode == "before":
-        patches[mode].extend(stmts)
+        patches[mode] = [*stmts, *patches[mode]] if mode in patches else stmts
     elif mode == "after":
-        patches[mode] = [*patches[mode], *stmts]
+        patches[mode] = [*patches[mode], *stmts] if mode in patches else stmts
     else:
         raise ValueError(f"Unknown patch mode: {mode!r}")
 
@@ -350,8 +344,8 @@ def __exit__(
 _cache_dir_lock = threading.Lock()
 
 
-def get_cache_dir() -> str:
-    """Get or create the temporary directory for awepatch."""
+def _user_cache_dir() -> str:
+    """Get or create the cache directory for awepatch."""
     global _cache_dir
 
     if _cache_dir is not None:
@@ -360,36 +354,41 @@ def get_cache_dir() -> str:
     with _cache_dir_lock:
         if _cache_dir is not None:
             return _cache_dir
+        from platformdirs import user_cache_dir
+
         _cache_dir = user_cache_dir("awepatch", appauthor=False, ensure_exists=True)
         return _cache_dir
 
 
-def write_patched_source(
+def persist_patched_source(
     source: str,
     name: str,
     type: str,
     origin: str = "",
 ) -> tuple[str, str]:
-    """Load a function's code object from its AST module.
+    """Persist the patched source code to a file and return the file path and source.
 
     Args:
-        source (str): The source code of the function.
-        name (str): The name of the function.
-        type (str): The type of the function (e.g., "module", "function").
-        origin (str, optional): The origin location for the function. Defaults to "".
+        source: The source code of the function.
+        name: The name of the function.
+        type: The type of the function (e.g., "module", "function").
+        origin: The origin location for the function. Defaults to "".
 
     Returns:
-        CodeType: The code object of the function.
+        tuple[str, str]: The file path and the source code of the function.
 
     """
+    from binascii import crc32
+
+    from filelock import FileLock
 
     origin = f" (patched from {origin})" if origin else ""
     source = f"# generated by awepatch{origin}\n{source}"
-
     bsource = source.encode("utf-8")
 
-    cache_dir = get_cache_dir()
+    cache_dir = _user_cache_dir()
     file_path = os.path.join(cache_dir, f"{type}_{name}_{crc32(bsource):010x}.py")
+
     with FileLock(f"{file_path}.lock"):
         if not os.path.exists(file_path):
             with open(file_path, "wb") as f: