Fix type annotations for profiler

lsst · Dec 13, 2023 · 451cdb3 · 451cdb3
1 parent fea2f20
commit 451cdb3
Showing 1 changed file with 7 additions and 2 deletions.
diff --git a/python/lsst/pipe/base/cmdLineTask.py b/python/lsst/pipe/base/cmdLineTask.py
@@ -25,17 +25,22 @@
 
 import contextlib
 import logging
+from collections.abc import Generator
+from typing import TYPE_CHECKING
 
 from deprecated.sphinx import deprecated
 
+if TYPE_CHECKING:
+    import cProfile
+
 
 @deprecated(
     reason="Replaced by lsst.utils.timer.profile().  Will be removed after v26.0",
     version="v25.0",
     category=FutureWarning,
 )
 @contextlib.contextmanager
-def profile(filename: str, log: logging.Logger | None = None):
+def profile(filename: str, log: logging.Logger | None = None) -> Generator[cProfile.Profile, None, None]:
     """Context manager for profiling with cProfile.
 
     Parameters
@@ -71,7 +76,7 @@ def profile(filename: str, log: logging.Logger | None = None):
     """
     if not filename:
         # Nothing to do
-        yield
+        yield None  # type: ignore
         return
     from cProfile import Profile