From f75853f9fe9921c5ed1aa4e6ab4cec594338905f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Mike=20M=C3=BCller?= <mmueller@python-academy.de>
Date: Mon, 25 Mar 2024 20:27:37 +0100
Subject: [PATCH] feat: add optional custom print callable (#2121)

---
 autotest/test_mbase.py | 31 +++++++++++++++++++++++++++++++
 flopy/mbase.py         | 11 +++++++++++
 flopy/mf6/mfsimbase.py |  7 +++++++
 3 files changed, 49 insertions(+)

diff --git a/autotest/test_mbase.py b/autotest/test_mbase.py
index 8e3ee5d828..2e8688cb3a 100644
--- a/autotest/test_mbase.py
+++ b/autotest/test_mbase.py
@@ -160,3 +160,34 @@ def test_run_model_exe_rel_path(mf6_model_path, function_tmpdir, use_ext):
         assert success
         assert any(buff)
         assert any(ws.glob("*.lst"))
+
+
+@pytest.mark.mf6
+@requires_exe("mf6")
+@pytest.mark.parametrize("use_paths", [True, False])
+@pytest.mark.parametrize(
+    "exe",
+    [
+        "mf6",
+        Path(which("mf6") or ""),
+        relpath_safe(Path(which("mf6") or "")),
+    ],
+)
+def test_run_model_custom_print(
+    mf6_model_path, function_tmpdir, use_paths, exe
+):
+    ws = function_tmpdir / "ws"
+    copytree(mf6_model_path, ws)
+
+    success, buff = run_model(
+        exe_name=exe,
+        namefile="mfsim.nam",
+        model_ws=ws if use_paths else str(ws),
+        silent=False,
+        report=True,
+        custom_print=print,
+    )
+
+    assert success
+    assert any(buff)
+    assert any(ws.glob("*.lst"))
diff --git a/flopy/mbase.py b/flopy/mbase.py
index c569944ff8..ee7bbd63fb 100644
--- a/flopy/mbase.py
+++ b/flopy/mbase.py
@@ -1744,6 +1744,7 @@ def run_model(
     normal_msg="normal termination",
     use_async=False,
     cargs=None,
+    custom_print=None,
 ) -> Tuple[bool, List[str]]:
     """
     Run the model using subprocess.Popen, optionally collecting stdout and printing
@@ -1782,12 +1783,22 @@ def run_model(
     cargs : str or list, optional, default None
         Additional command line arguments to pass to the executable.
         (Default is None)
+    custom_print: callable
+        Optional callable for printing. It will replace the builtin print
+        function. This is useful for a shorter print output or integration into
+        other systems such as GUIs.
+        default is None, i.e. use the builtin print
     Returns
     -------
     success : boolean
     buff : list of lines of stdout (empty if report is False)
 
     """
+    if custom_print is not None:
+        print = custom_print
+    else:
+        print = __builtins__["print"]
+
     success = False
     buff = []
 
diff --git a/flopy/mf6/mfsimbase.py b/flopy/mf6/mfsimbase.py
index e84c4dfc51..163eab0d9c 100644
--- a/flopy/mf6/mfsimbase.py
+++ b/flopy/mf6/mfsimbase.py
@@ -1631,6 +1631,7 @@ def run_simulation(
         normal_msg="normal termination",
         use_async=False,
         cargs=None,
+        custom_print=None,
     ):
         """
         Run the simulation.
@@ -1657,6 +1658,11 @@ def run_simulation(
             cargs : str or list of strings
                 Additional command line arguments to pass to the executable.
                 default is None
+            custom_print: callable
+                Optional callbale for printing. It will replace the builtin
+                print function. This is useful for shorter prints or integration
+                into other systems such as GUIs.
+                default is None, i.e. use the builtion print
 
         Returns
         --------
@@ -1683,6 +1689,7 @@ def run_simulation(
             normal_msg=normal_msg,
             use_async=use_async,
             cargs=cargs,
+            custom_print=custom_print,
         )
 
     def delete_output_files(self):