Import memory_profiler only when it is enabled

Belval · Jan 7, 2024 · e8a9235 · e8a9235
1 parent 351ec66
commit e8a9235
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/tests.py b/tests.py
@@ -11,7 +11,6 @@
 from subprocess import Popen, PIPE
 from tempfile import TemporaryDirectory
 from multiprocessing.dummy import Pool
-from memory_profiler import profile as profile_memory
 
 sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
 
@@ -31,6 +30,8 @@
 from functools import wraps
 
 PROFILE_MEMORY = os.environ.get("PROFILE_MEMORY", False)
+if PROFILE_MEMORY:
+    from memory_profiler import profile as profile_memory
 
 try:
     subprocess.call(