From 0df56117e4730d8dbb20fb06d6f5a1c186c1355f Mon Sep 17 00:00:00 2001
From: Philipp A <flying-sheep@web.de>
Date: Mon, 5 May 2025 11:00:31 +0200
Subject: [PATCH] Backport PR #3620: Switch from profimp to tuna for parsing
 import profiles

---
 pyproject.toml            |  2 +-
 tests/test_performance.py | 48 ++++++++++++++++++++++++++++++---------
 2 files changed, 38 insertions(+), 12 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index baf05cc91a..26fdc719e6 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -90,7 +90,7 @@ test-min = [
     "pytest-xdist[psutil]",
     "pytest-randomly",
     "pytest-rerunfailures",
-    "profimp",
+    "tuna",
 ]
 test = [
     "scanpy[test-min]",
diff --git a/tests/test_performance.py b/tests/test_performance.py
index 26efe44cf7..eb6b042a02 100644
--- a/tests/test_performance.py
+++ b/tests/test_performance.py
@@ -1,39 +1,65 @@
 from __future__ import annotations
 
-import json
+import enum
 import sys
+from pathlib import Path
 from subprocess import run
+from tempfile import NamedTemporaryFile
+from typing import TYPE_CHECKING, TypedDict, cast
 
+if TYPE_CHECKING:
+    from collections.abc import Iterable, MutableSet
+    from typing import NotRequired
 
-def descend(profimp_data, modules, path):
-    module = profimp_data["module"]
+
+class TunaColor(enum.IntEnum):
+    Func = 0
+    Builtin = 1
+    Deprecated = 2
+
+
+class TunaProf(TypedDict):
+    text: list[str]
+    value: float
+    color: TunaColor
+    children: NotRequired[list[TunaProf]]
+
+
+def descend(
+    profile: TunaProf, modules: MutableSet[str], path: Iterable[str] = ()
+) -> Iterable[str]:
+    [module] = profile["text"]
     path = [*path, module]
     if module in modules:
         yield " → ".join(e for e in path if e is not None)
         modules.remove(module)
-    for child in profimp_data["children"]:
+    for child in profile.get("children", []):
         yield from descend(child, modules, path)
 
 
-def get_import_paths(modules):
+def get_import_paths(modules: Iterable[str]) -> Iterable[str]:
+    from tuna import read_import_profile
+
     proc = run(
-        [sys.executable, "-m", "profimp.main", "import scanpy"],
+        [sys.executable, "-X", "importtime", "-c", "import scanpy"],
         capture_output=True,
         check=True,
     )
-    data = json.loads(proc.stdout)
-    return descend(data, set(modules), [])
+    with NamedTemporaryFile() as f:
+        Path(f.name).write_bytes(proc.stderr)
+        data = cast("TunaProf", read_import_profile(f.name))
+    return descend(data, set(modules))
 
 
-def test_deferred_imports(imported_modules):
+def test_deferred_imports(imported_modules: frozenset[str]) -> None:
     slow_to_import = {
         "umap",  # neighbors, tl.umap
         "seaborn",  # plotting
         "sklearn.metrics",  # neighbors
         "pynndescent",  # neighbors
         "networkx",  # diffmap, paga, plotting._utils
-        # TODO: 'matplotlib.pyplot',
-        # TODO (maybe): 'numba',
+        # TODO: "matplotlib.pyplot",
+        # TODO (maybe): "numba",
     }
     falsely_imported = slow_to_import & imported_modules