Add tests and fix bugs

Author: datapythonista

doc/source/whatsnew/v3.0.0.rst

@@ -88,6 +88,7 @@ Other enhancements
 - Support passing a :class:`Iterable[Hashable]` input to :meth:`DataFrame.drop_duplicates` (:issue:`59237`)
 - Support reading Stata 102-format (Stata 1) dta files (:issue:`58978`)
 - Support reading Stata 110-format (Stata 7) dta files (:issue:`47176`)
+- Third-party packages can now register engines that can be used in pandas I/O operations :func:`read_iceberg` and :meth:`DataFrame.to_iceberg` (:issue:`61584`)
 
 .. ---------------------------------------------------------------------------
 .. _whatsnew_300.notable_bug_fixes:

pandas/core/frame.py

@@ -188,7 +188,10 @@
     nargsort,
 )
 
-from pandas.io.common import get_handle
+from pandas.io.common import (
+    allow_third_party_engines,
+    get_handle,
+)
 from pandas.io.formats import (
     console,
     format as fmt,
@@ -3547,6 +3550,7 @@ def to_xml(
 
         return xml_formatter.write_output()
 
+    @allow_third_party_engines
     def to_iceberg(
         self,
         table_identifier: str,
@@ -3556,6 +3560,7 @@ def to_iceberg(
         location: str | None = None,
         append: bool = False,
         snapshot_properties: dict[str, str] | None = None,
+        engine: str | None = None,
     ) -> None:
         """
         Write a DataFrame to an Apache Iceberg table.
@@ -3580,6 +3585,10 @@ def to_iceberg(
             If ``True``, append data to the table, instead of replacing the content.
         snapshot_properties : dict of {str: str}, optional
             Custom properties to be added to the snapshot summary
+        engine : str, optional
+            The engine to use. Engines can be installed via third-party packages. For an
+            updated list of existing pandas I/O engines check the I/O engines section of
+            our Ecosystem page.
 
         See Also
         --------

pandas/io/common.py

@@ -9,13 +9,15 @@
 import codecs
 from collections import defaultdict
 from collections.abc import (
+    Callable,
     Hashable,
     Mapping,
     Sequence,
 )
 import dataclasses
 import functools
 import gzip
+from importlib.metadata import entry_points
 from io import (
     BufferedIOBase,
     BytesIO,
@@ -51,8 +53,6 @@
 import warnings
 import zipfile
 
-import pkg_resources
-
 from pandas._typing import (
     BaseBuffer,
     ReadCsvBuffer,
@@ -1290,9 +1290,9 @@ def dedup_names(
     return names
 
 
-def _engine_func(format_name: str, engine_name: str, is_writer: bool):
+def _get_io_engine(name: str):
     """
-    Return the engine function for a given format and operation.
+    Return an I/O engine by its name.
 
     pandas I/O engines can be registered via entry points. The first time this
     function is called it will register all the entry points of the "pandas.io_engine"
@@ -1304,13 +1304,8 @@ def _engine_func(format_name: str, engine_name: str, is_writer: bool):
 
     Parameters
     ----------
-    format_name : str
-        The format such as 'csv', 'parquet', 'json', 'html', etc.
-    engine_name : str
+    name : str
         The engine name provided by the user in `engine=<value>`.
-    is_writer : bool
-        `True` to return the `to_<format>` function, `False` to return the
-        `read_<format>` one.
 
     Examples
     --------
@@ -1330,59 +1325,57 @@ def _engine_func(format_name: str, engine_name: str, is_writer: bool):
 
     ```
 
-    Then the `read_csv` method of the engine can be retrieved with:
+    Then the `read_csv` method of the engine can be used with:
 
-    >>> func = _engine_func(format_name="csv", engine_name="dummy", is_writer=False)
+    >>> _get_io_engine(engine_name="dummy").read_csv("myfile.csv")  # doctest: +SKIP
 
     This is used internally to dispatch the next pandas call to the engine caller:
 
-    >>> df = read_csv("myfile.csv", engine="dummy")
+    >>> df = read_csv("myfile.csv", engine="dummy")  # doctest: +SKIP
     """
     global _io_engines
 
     if _io_engines is None:
         _io_engines = {}
-        for entry_point in pkg_resources.iter_entry_points(group="pandas.io_engine"):
-            _io_engines[entry_point.name] = entry_point.load()
+        for entry_point in entry_points().select(group="pandas.io_engine"):
+            package_name = entry_point.dist.metadata["Name"]
+            if entry_point.name in _io_engines:
+                _io_engines[entry_point.name]._other_providers.append(package_name)
+            else:
+                _io_engines[entry_point.name] = entry_point.load()
+                _io_engines[entry_point.name]._provider_name = package_name
+                _io_engines[entry_point.name]._other_providers = []
 
     try:
-        engine_class = _io_engines[engine_name]
+        engine = _io_engines[name]
     except KeyError as err:
         raise ValueError(
-            f"'{engine_name}' is not a known engine. Some engines are only available "
+            f"'{name}' is not a known engine. Some engines are only available "
             "after installing the package that provides them."
         ) from err
 
-    func_name = f"to_{format_name}" if is_writer else f"read_{format_name}"
-    try:
-        engine_method = getattr(engine_class, func_name)
-    except AttributeError as err:
-        raise ValueError(
-            f"The engine '{engine_name}' does not provide a '{func_name}' function"
-        ) from err
-    else:
-        return engine_method
-
-
-def _extract_io_function_info(func_name):
-    """
-    Return the format and if it's a reader or writer from a function name like read_csv.
-    """
-    op_type, format_name = func_name.split("_", maxsplit=1)
-    if op_type == "read":
-        is_writer = False
-    elif op_type == "to":
-        is_writer = True
-    else:
-        raise ValueError(
-            "Unable to extract info from the function name '{func_name}'. "
-            "The expected format is `read_<format> or `to_<format>`."
+    if engine._other_providers:
+        msg = (
+            f"The engine '{name}' has been registered by the package "
+            f"'{engine._provider_name}' and will be used. "
         )
+        if len(engine._other_providers):
+            msg += (
+                "The package '{engine._other_providers}' also tried to register "
+                "the engine, but it couldn't because it was already registered."
+            )
+        else:
+            msg += (
+                "Other packages that tried to register the engine, but they couldn't "
+                "because it was already registered are: "
+                f"{str(engine._other_providers)[1:-1]}."
+            )
+        warnings.warn(RuntimeWarning, msg, stacklevel=find_stack_level())
 
-    return format_name, is_writer
+    return engine
 
 
-def allow_third_party_engines(skip_engines: list[str] | None = None):
+def allow_third_party_engines(skip_engines: list[str] | Callable | None = None):
     """
     Decorator to avoid boilerplate code when allowing readers and writers to use
     third-party engines.
@@ -1415,14 +1408,21 @@ def allow_third_party_engines(skip_engines: list[str] | None = None):
     def decorator(func):
         @functools.wraps(func)
         def wrapper(*args, **kwargs):
-            if "engine" in kwargs and kwargs["engine"] not in skip_engines:
-                format_name, is_writer = _extract_io_function_info(func.__name__)
-                engine_func = _engine_func(
-                    format_name=format_name,
-                    engine_name=kwargs.pop("engine"),
-                    is_writer=is_writer,
-                )
-                return engine_func(*args, **kwargs)
+            if callable(skip_engines):
+                skip_engine = False
+            else:
+                skip_engine = kwargs["engine"] in skip_engines
+
+            if "engine" in kwargs and not skip_engine:
+                engine_name = kwargs.pop("engine")
+                engine = _get_io_engine(engine_name)
+                try:
+                    return getattr(engine, func.__name__)(*args, **kwargs)
+                except AttributeError as err:
+                    raise ValueError(
+                        f"The engine '{engine_name}' does not provide a "
+                        f"'{func.__name__}' function"
+                    ) from err
             else:
                 return func(*args, **kwargs)
 

pandas/io/iceberg.py

@@ -9,7 +9,7 @@
 from pandas.io.common import allow_third_party_engines
 
 
-@allow_third_party_engines()
+@allow_third_party_engines
 def read_iceberg(
     table_identifier: str,
     catalog_name: str | None = None,

pandas/tests/io/test_io_engines.py

@@ -0,0 +1,48 @@
+import pytest
+
+from pandas.io import common
+
+
+@pytest.fixture
+def patch_engine(monkeypatch):
+    class MockIoEngine:
+        @classmethod
+        def read_foo(cls, fname):
+            return "third-party"
+
+    monkeypatch.setattr(common, "_get_io_engine", lambda name: MockIoEngine)
+
+
+class TestIoEngines:
+    def test_decorator_with_no_engine(self, patch_engine):
+        @common.allow_third_party_engines
+        def read_foo(fname, engine=None):
+            return "default"
+
+        result = read_foo("myfile.foo")
+        assert result == "default"
+
+    def test_decorator_with_skipped_engine(self, patch_engine):
+        @common.allow_third_party_engines(skip_engines=["c"])
+        def read_foo(fname, engine=None):
+            return "default"
+
+        result = read_foo("myfile.foo", engine="c")
+        assert result == "default"
+
+    def test_decorator_with_third_party_engine(self, patch_engine):
+        @common.allow_third_party_engines
+        def read_foo(fname, engine=None):
+            return "default"
+
+        result = read_foo("myfile.foo", engine="third-party")
+        assert result == "third-party"
+
+    def test_decorator_with_third_party_engine_but_no_method(self, patch_engine):
+        @common.allow_third_party_engines
+        def read_bar(fname, engine=None):
+            return "default"
+
+        msg = "'third-party' does not provide a 'read_bar'"
+        with pytest.raises(ValueError, match=msg):
+            read_bar("myfile.foo", engine="third-party")