From 1dded5cc9b60f792a73673555fdac878b3f15f82 Mon Sep 17 00:00:00 2001
From: papertager <2567587994@qq.com>
Date: Tue, 9 Jun 2026 21:22:41 +0800
Subject: [PATCH] Allow profiler output directory override

---
 mcoplib/profiler.py                       |  8 +++++-
 unit_test/test_profiler_output_dir_env.py | 33 +++++++++++++++++++++++
 2 files changed, 40 insertions(+), 1 deletion(-)
 create mode 100644 unit_test/test_profiler_output_dir_env.py

diff --git a/mcoplib/profiler.py b/mcoplib/profiler.py
index adfba41..e8ffa05 100644
--- a/mcoplib/profiler.py
+++ b/mcoplib/profiler.py
@@ -19,6 +19,10 @@ def _timestamp() -> str:
     return datetime.now().strftime("%Y%m%dT%H%M%S")
 
 
+def _resolve_output_dir(default_output_dir):
+    return os.getenv("PROFILER_OUTPUT_DIR", default_output_dir)
+
+
 def _track_handler(prof, output_dir, func_name):
     """
     Track handler implementation that matches test.py track_handler format.
@@ -120,7 +124,9 @@ def wrapper(*args, **kwargs):
                             active=1,
                             repeat=repeat
                         ),
-                        on_trace_ready=lambda prof: _track_handler(prof, output_dir, func.__name__),
+                on_trace_ready=lambda prof: _track_handler(
+                    prof, _resolve_output_dir(output_dir), func.__name__
+                ),
                         with_modules=True,
                         record_shapes=True,
                         profile_memory=True
diff --git a/unit_test/test_profiler_output_dir_env.py b/unit_test/test_profiler_output_dir_env.py
new file mode 100644
index 0000000..3d0d70a
--- /dev/null
+++ b/unit_test/test_profiler_output_dir_env.py
@@ -0,0 +1,33 @@
+import os
+import sys
+import unittest
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+
+from mcoplib.profiler import _resolve_output_dir
+
+
+class ProfilerOutputDirEnvTest(unittest.TestCase):
+    def test_uses_default_without_env(self):
+        old = os.environ.pop("PROFILER_OUTPUT_DIR", None)
+        try:
+            self.assertEqual(_resolve_output_dir("./profiles"), "./profiles")
+        finally:
+            if old is not None:
+                os.environ["PROFILER_OUTPUT_DIR"] = old
+
+    def test_env_overrides_default(self):
+        old = os.environ.get("PROFILER_OUTPUT_DIR")
+        os.environ["PROFILER_OUTPUT_DIR"] = "/tmp/mcoplib-profiles"
+        try:
+            self.assertEqual(_resolve_output_dir("./profiles"), "/tmp/mcoplib-profiles")
+        finally:
+            if old is None:
+                os.environ.pop("PROFILER_OUTPUT_DIR", None)
+            else:
+                os.environ["PROFILER_OUTPUT_DIR"] = old
+
+
+if __name__ == "__main__":
+    unittest.main()