You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by ti...@apache.org on 2026/09/04 20:09:30 UTC
(datafusion-python) 02/02: Tidy the loose ends from review of with_extensions
This is an automated email from the ASF dual-hosted git repository.
timsaucer pushed a commit to branch feat/ffi-with-extensions
in repository https://gitbox.apache.org/repos/asf/datafusion-python.git
commit e06662acb7007c5493dea5f457fdef26108ecd50
Author: Tim Saucer <ti...@gmail.com>
AuthorDate: Fri Sep 4 16:07:49 2026 -0400
Tidy the loose ends from review of with_extensions
Mark `SessionExtensionExportable` `@runtime_checkable` and have
`with_extensions` check it with `isinstance` rather than `hasattr`, so the
annotation and the runtime check are the same statement, and callers can ask
the question too. Covered by a doctest on the protocol.
Replace the leading-underscore skip in `test_wrapper_coverage` with a named
allowlist. The pattern also excused `DataFrame._repr_html_`, which a wrapper
does have to provide, so a two-method need was weakening coverage for every
private name. Removing `_install_extensions` from the allowlist fails the test,
so the entry is load-bearing rather than decorative.
Say in `_CodecOnlyExtension` that retaining the context is what the protocol
tells real extensions not to do, and that it is kept only so a test can assert
which context the factory was handed.
Let the docstring-example shim in the query planner example accept a config
positionally, the way the real constructor does. Editing the docstring to
`SessionContext(config)` now fails as a doctest diff rather than as a
`TypeError` inside the harness.
Co-Authored-By: Claude Opus 5 (1M context) <no...@anthropic.com>
---
.../tests/_test_three_library_query_planner.py | 10 ++++++----
python/datafusion/context.py | 2 +-
python/datafusion/extensions.py | 21 +++++++++++++++++++-
python/tests/test_context.py | 8 +++++++-
python/tests/test_wrapper_coverage.py | 23 ++++++++++++++++------
5 files changed, 51 insertions(+), 13 deletions(-)
diff --git a/examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py b/examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py
index 4ee9b393..ed9e6184 100644
--- a/examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py
+++ b/examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py
@@ -1099,10 +1099,12 @@ def test_with_extensions_docstring_example_still_runs():
module = types.ModuleType("my_extension")
module.DistributedEngineExtension = _DocstringExampleExtension
- def make_context() -> SessionContext:
- return SessionContext(
- SessionConfig().with_extension(MyPlannerConfig(max_rows=3))
- )
+ def make_context(config: SessionConfig | None = None) -> SessionContext:
+ # Accept a config so the example is free to pass one. Supplying it
+ # positionally the way the real constructor does keeps a docstring
+ # edit failing as a doctest diff rather than as a TypeError in here.
+ config = SessionConfig() if config is None else config
+ return SessionContext(config.with_extension(MyPlannerConfig(max_rows=3)))
test = doctest.DocTest(
examples,
diff --git a/python/datafusion/context.py b/python/datafusion/context.py
index 1ea937ac..ac453bb7 100644
--- a/python/datafusion/context.py
+++ b/python/datafusion/context.py
@@ -1899,7 +1899,7 @@ class SessionContext:
msg = "with_extensions requires at least one extension"
raise ValueError(msg)
for extension in extensions:
- if not hasattr(extension, "__datafusion_session_extension__"):
+ if not isinstance(extension, SessionExtensionExportable):
msg = (
"Extension does not implement __datafusion_session_extension__: "
f"{extension!r}"
diff --git a/python/datafusion/extensions.py b/python/datafusion/extensions.py
index c12e631c..8e35c720 100644
--- a/python/datafusion/extensions.py
+++ b/python/datafusion/extensions.py
@@ -37,7 +37,7 @@ extensions guide in the contributor documentation for the full rationale.
from __future__ import annotations
from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any, Protocol
+from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable
if TYPE_CHECKING:
from _typeshed import CapsuleType as _PyCapsule
@@ -124,9 +124,15 @@ class SessionExtensionComponents:
"""
+@runtime_checkable
class SessionExtensionExportable(Protocol):
"""Type hint for extension bundles installable via ``with_extensions``.
+ Runtime-checkable, so ``isinstance`` answers whether an object implements
+ the protocol. Only the presence of the method is checked, which is the same
+ question :py:meth:`~datafusion.context.SessionContext.with_extensions` asks
+ before calling it.
+
Implementations are reusable configuration objects: they must create fresh
components on every call using the context supplied by
:py:meth:`~datafusion.context.SessionContext.with_extensions`, and must not
@@ -134,6 +140,19 @@ class SessionExtensionExportable(Protocol):
next call may install onto a different session. They should also avoid
mutating the context they are handed — a registration made during binding
is not rolled back if a later extension fails.
+
+ Examples:
+ >>> from datafusion import (
+ ... SessionExtensionComponents,
+ ... SessionExtensionExportable,
+ ... )
+ >>> class MyLibraryExtension:
+ ... def __datafusion_session_extension__(self, ctx):
+ ... return SessionExtensionComponents()
+ >>> isinstance(MyLibraryExtension(), SessionExtensionExportable)
+ True
+ >>> isinstance(object(), SessionExtensionExportable)
+ False
"""
def __datafusion_session_extension__( # noqa: D105
diff --git a/python/tests/test_context.py b/python/tests/test_context.py
index ff53c177..0418acde 100644
--- a/python/tests/test_context.py
+++ b/python/tests/test_context.py
@@ -882,7 +882,13 @@ def test_contexts_sharing_a_session_share_the_planner(ctx):
class _CodecOnlyExtension:
- """Contributes decline-all codecs exported from an unrelated session."""
+ """Contributes decline-all codecs exported from an unrelated session.
+
+ Retaining ``ctx`` is what the protocol tells real extensions not to do —
+ a bundle is reusable, so a cached context belongs to whichever session it
+ was last installed on. It is kept here only so a test can assert *which*
+ context the factory was handed.
+ """
def __init__(self):
self.exporter = SessionContext()
diff --git a/python/tests/test_wrapper_coverage.py b/python/tests/test_wrapper_coverage.py
index b1afd683..7e8bd3f2 100644
--- a/python/tests/test_wrapper_coverage.py
+++ b/python/tests/test_wrapper_coverage.py
@@ -28,6 +28,17 @@ except ImportError:
from enum import EnumMeta as EnumType
+# Internal methods a wrapper calls but does not re-export. Add to this only
+# when the method exists to serve a public wrapper, never to silence a genuine
+# gap in coverage.
+PRIVATE_SUPPORT_METHODS = frozenset(
+ {
+ # Support method for SessionContext.with_extensions.
+ "_install_extensions",
+ }
+)
+
+
def _check_enum_exports(internal_obj, wrapped_obj) -> None:
"""Check that all enum values are present in wrapped object."""
expected_values = [v for v in dir(internal_obj) if not v.startswith("__")]
@@ -67,12 +78,12 @@ def missing_exports(internal_obj, wrapped_obj) -> None:
pytest.fail(f"Missing __repr__: {internal_obj.__name__}")
for internal_attr_name in dir(internal_obj):
- # Single-underscore names are private support methods for the
- # wrappers (e.g. SessionContext._install_extensions) and are not
- # part of the public surface that requires a wrapper.
- if internal_attr_name.startswith("_") and not internal_attr_name.startswith(
- "__"
- ):
+ # Private support methods that exist only for a wrapper to call, so
+ # they are not part of the public surface and need no wrapper of their
+ # own. Listed rather than matched by leading underscore, which would
+ # also excuse names like `_repr_html_` that a wrapper does have to
+ # provide.
+ if internal_attr_name in PRIVATE_SUPPORT_METHODS:
continue
wrapped_attr_name = internal_attr_name.removeprefix("Raw")
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org