From 2159288aedc47405db5011f05664ea9fb5e1ff8b Mon Sep 17 00:00:00 2001
From: Pierre-Marie de Rodat <derodat@adacore.com>
Date: Fri, 23 Aug 2024 09:57:02 +0200
Subject: [PATCH] PatternSubstitute: update annotations to accept replacement
 callbacks

---
 NEWS.md                         |  1 +
 src/e3/testsuite/driver/diff.py | 21 ++++++++++++++++++---
 tests/tests/test_diff.py        | 11 +++++++++++
 3 files changed, 30 insertions(+), 3 deletions(-)

diff --git a/NEWS.md b/NEWS.md
index 98d3972..1b25be4 100644
--- a/NEWS.md
+++ b/NEWS.md
@@ -1,6 +1,7 @@
 27.0 (Not released yet)
 =======================
 
+* `PatternSubstitute`: update annotations to accept replacement callbacks.
 * `e3.testsuite.report.xunit`: add a `XUnitImportApp` class to make it possible
   to create customized versions of the `e3-convert-xunit` script without code
   duplication.
diff --git a/src/e3/testsuite/driver/diff.py b/src/e3/testsuite/driver/diff.py
index 9e590e2..c8a08c4 100644
--- a/src/e3/testsuite/driver/diff.py
+++ b/src/e3/testsuite/driver/diff.py
@@ -1,6 +1,17 @@
+from __future__ import annotations
+
 import os
 import re
-from typing import AnyStr, Generic, List, Optional, Pattern, Tuple, Union
+from typing import (
+    AnyStr,
+    Callable,
+    Generic,
+    List,
+    Optional,
+    Pattern,
+    Tuple,
+    Union,
+)
 
 from e3.diff import diff
 from e3.os.fs import unixpath
@@ -98,7 +109,9 @@ class PatternSubstitute(OutputRefiner, Generic[AnyStr]):
     """Replace patterns in outputs."""
 
     def __init__(
-        self, pattern: AnyStr, replacement: Optional[AnyStr] = None
+        self,
+        pattern: AnyStr,
+        replacement: AnyStr | Callable[[re.Match], AnyStr] | None = None,
     ) -> None:
         """
         Initialize a PatternSubstitute instance.
@@ -109,7 +122,9 @@ def __init__(
             string).
         """
         self.regexp: Pattern[AnyStr] = re.compile(pattern)
-        self.replacement: AnyStr = replacement or type(pattern)()
+        self.replacement: AnyStr | Callable[[re.Match], AnyStr] = (
+            replacement or type(pattern)()
+        )
 
     def refine(self, output: AnyStr) -> AnyStr:
         return self.regexp.sub(self.replacement, output)
diff --git a/tests/tests/test_diff.py b/tests/tests/test_diff.py
index 2551ec2..e5e6b29 100644
--- a/tests/tests/test_diff.py
+++ b/tests/tests/test_diff.py
@@ -361,3 +361,14 @@ def test_line_by_line():
 
     assert subst.refine(output) == b"first some-content line\n"
     assert lbl_subst.refine(output) == b"first line\nsome-content line\n"
+
+
+def test_pattern_substitute_callback():
+    """Check that PatternSubstute accept replacement callbacks."""
+
+    def repl(m):
+        return m.group(1)
+
+    refiner = diff.PatternSubstitute(r"[a-z]*\((.*)\)", repl)
+
+    assert refiner.refine("foo(1, 2)") == "1, 2"