ExperimentAxisQuery uses the thread pool from ContextBase

python-spec/src/somacore/__init__.py

@@ -30,6 +30,7 @@
 from .query import AxisColumnNames
 from .query import AxisQuery
 from .query import ExperimentAxisQuery
+from .types import ContextBase
 
 try:
     # This trips up mypy since it's a generated file:
@@ -59,4 +60,5 @@
     "AxisColumnNames",
     "AxisQuery",
     "ExperimentAxisQuery",
+    "ContextBase",
 )

python-spec/src/somacore/query/query.py

@@ -25,6 +25,7 @@
 from .. import data
 from .. import measurement
 from .. import options
+from ..types import ContextBase
 from . import _fast_csr
 from . import axis
 from . import types
@@ -587,11 +588,13 @@ def _var_df(self) -> data.DataFrame:
 
     @property
     def _threadpool(self) -> futures.ThreadPoolExecutor:
-        """Creates a thread pool just in time."""
+        """
+        Returns the threadpool provided by the experiment's context.
+        If not available, creates a thread pool just in time."""
+        if self.experiment.context._threadpool is not None:
+            return self.experiment.context._threadpool
+
         if self._threadpool_ is None:
-            # TODO: the user should be able to set their own threadpool, a la asyncio's
-            # loop.set_default_executor().  This is important for managing the level of
-            # concurrency, etc.
             self._threadpool_ = futures.ThreadPoolExecutor()
         return self._threadpool_
 
@@ -797,6 +800,10 @@ def ms(self) -> Mapping[str, measurement.Measurement]:
     def obs(self) -> data.DataFrame:
         ...
 
+    @property
+    def context(self) -> ContextBase:
+        ...
+
 
 class _HasObsVar(Protocol[_T_co]):
     """Something which has an ``obs`` and ``var`` field.

python-spec/src/somacore/types.py

@@ -6,6 +6,7 @@
 """
 
 import sys
+from concurrent import futures
 from typing import TYPE_CHECKING, NoReturn, Optional, Sequence, Type, TypeVar
 
 from typing_extensions import Protocol, TypeGuard
@@ -75,3 +76,12 @@ def is_slice_of(__obj: object, __typ: Type[_T]) -> TypeGuard[Slice[_T]]:
         and (__obj.stop is None or isinstance(__obj.stop, __typ))
         and (__obj.step is None or isinstance(__obj.step, __typ))
     )
+
+
+class ContextBase(Protocol):
+    """A protocol for a context manager that can be used as a base class.
+
+    The only requirement for somacore is that it should contain a threadpool.
+    """
+
+    _threadpool: futures.ThreadPoolExecutor