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) <noreply@anthropic.com>
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 4ee9b39..ed9e618 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 @@ 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 1ea937a..ac453bb 100644 --- a/python/datafusion/context.py +++ b/python/datafusion/context.py
@@ -1899,7 +1899,7 @@ 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 c12e631..8e35c72 100644 --- a/python/datafusion/extensions.py +++ b/python/datafusion/extensions.py
@@ -37,7 +37,7 @@ 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 @@ """ +@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 @@ 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 ff53c17..0418acd 100644 --- a/python/tests/test_context.py +++ b/python/tests/test_context.py
@@ -882,7 +882,13 @@ 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 b1afd68..7e8bd3f 100644 --- a/python/tests/test_wrapper_coverage.py +++ b/python/tests/test_wrapper_coverage.py
@@ -28,6 +28,17 @@ 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 @@ 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")