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/base.py

@@ -10,6 +10,7 @@
 from typing_extensions import LiteralString, Self
 
 from . import options
+from .types import ContextBase
 
 
 class SOMAObject(metaclass=abc.ABCMeta):
@@ -24,7 +25,7 @@ def open(
         uri: str,
         mode: options.OpenMode = "r",
         *,
-        context: Optional[Any] = None,
+        context: Optional[ContextBase] = None,
         platform_config: Optional[options.PlatformConfig] = None,
     ) -> Self:
         """Opens the SOMA object of this type at the given URI.
@@ -42,7 +43,7 @@ def open(
 
     @classmethod
     @abc.abstractmethod
-    def exists(cls, uri: str, *, context: Optional[Any] = None) -> bool:
+    def exists(cls, uri: str, *, context: Optional[ContextBase] = None) -> bool:
         """Checks whether a SOMA object of this type is stored at the URI.
 
         Args:
@@ -65,7 +66,7 @@ def uri(self) -> str:
         raise NotImplementedError()
 
     @property
-    def context(self) -> Any:
+    def context(self) -> Optional[ContextBase]:
         """A value storing implementation-specific configuration information.
 
         This contains long-lived (i.e., not call-specific) information that is

python-spec/src/somacore/collection.py

@@ -1,12 +1,13 @@
 import abc
-from typing import Any, MutableMapping, Optional, Sequence, Type, TypeVar, overload
+from typing import MutableMapping, Optional, Sequence, Type, TypeVar, overload
 
 import pyarrow as pa
 from typing_extensions import Final, Self
 
 from . import base
 from . import data
 from . import options
+from .types import ContextBase
 
 _Elem = TypeVar("_Elem", bound=base.SOMAObject)
 """Element Type for a SOMA collection."""
@@ -35,7 +36,7 @@ def create(
         uri: str,
         *,
         platform_config: Optional[options.PlatformConfig] = None,
-        context: Optional[Any] = None,
+        context: Optional[ContextBase] = None,
     ) -> Self:
         """Creates a new collection of this type at the given URI.
 

python-spec/src/somacore/data.py

@@ -23,6 +23,7 @@
 
 from . import base
 from . import options
+from .types import ContextBase
 
 _RO_AUTO = options.ResultOrder.AUTO
 
@@ -48,7 +49,7 @@ def create(
         index_column_names: Sequence[str] = (options.SOMA_JOINID,),
         domain: Optional[Sequence[Optional[Tuple[Any, Any]]]] = None,
         platform_config: Optional[options.PlatformConfig] = None,
-        context: Optional[Any] = None,
+        context: Optional[ContextBase] = None,
     ) -> Self:
         """Creates a new ``DataFrame`` at the given URI.
 
@@ -231,7 +232,7 @@ def create(
         type: pa.DataType,
         shape: Sequence[Optional[int]],
         platform_config: Optional[options.PlatformConfig] = None,
-        context: Optional[Any] = None,
+        context: Optional[ContextBase] = None,
     ) -> Self:
         """Creates a new ND array of the current type at the given URI.
 

python-spec/src/somacore/ephemeral/collections.py

@@ -8,6 +8,7 @@
 from .. import experiment
 from .. import measurement
 from .. import options
+from ..types import ContextBase
 
 _Elem = TypeVar("_Elem", bound=base.SOMAObject)
 
@@ -51,7 +52,9 @@ def open(cls, *args, **kwargs) -> NoReturn:
         )
 
     @classmethod
-    def exists(cls, uri: str, *, context: Any = None) -> Literal[False]:
+    def exists(
+        cls, uri: str, *, context: Optional[ContextBase] = None
+    ) -> Literal[False]:
         del uri, context  # All unused.
         # Ephemeral collections are in-memory only and do not otherwise exist.
         return False

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,16 @@ 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 is not None
+            and 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 +803,10 @@ def ms(self) -> Mapping[str, measurement.Measurement]:
     def obs(self) -> data.DataFrame:
         ...
 
+    @property
+    def context(self) -> Optional[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.
+    If a threadpool is specified as part of the context, it will be used by the
+    implementer. Otherwise, the implementer will use its own threadpool.
+    """
+
+    threadpool: Optional[futures.ThreadPoolExecutor]