Fugue plugin (#311)

* plugin

* update

RELEASE.md

@@ -7,6 +7,7 @@
 -   A [smarter](https://github.com/fugue-project/fugue/issues/304) way to determine default npartitions in Dask
 -   Support [even partitioning](https://github.com/fugue-project/fugue/issues/303) on Dask
 -   Add handling of [nested ArrayType](https://github.com/fugue-project/fugue/issues/308) on Spark
+-   Change to [plugin approach](https://github.com/fugue-project/fugue/issues/310) to avoid explicit import
 
 ## 0.6.5
 

fugue/_utils/interfaceless.py

@@ -263,6 +263,8 @@ def _parse_param(  # noqa: C901
         param: Optional[inspect.Parameter],
         none_as_other: bool = True,
     ) -> "_FuncParam":
+        import fugue._utils.register  # pylint: disable=W0611 # noqa: F401
+
         if annotation == type(None):  # noqa: E721
             return _NoneParam(param)
         if annotation == inspect.Parameter.empty:

fugue/_utils/register.py

@@ -0,0 +1,17 @@
+try:
+    from importlib.metadata import entry_points  # type:ignore
+except ImportError:  # pragma: no cover
+    from importlib_metadata import entry_points  # type:ignore
+
+
+def register_plugins():
+    for plugin in entry_points().get("fugue.plugins", []):
+        try:
+            register_func = plugin.load()
+            assert callable(register_func), f"{plugin.name} is not a callable"
+            register_func()
+        except ImportError:  # pragma: no cover
+            pass
+
+
+register_plugins()

fugue/execution/__init__.py

@@ -14,16 +14,21 @@
     SqliteEngine,
 )
 
-register_execution_engine(
-    "native", lambda conf: NativeExecutionEngine(conf), on_dup="ignore"
-)
-register_execution_engine(
-    "pandas", lambda conf: NativeExecutionEngine(conf), on_dup="ignore"
-)
-register_sql_engine("sqlite", lambda engine: SqliteEngine(engine), on_dup="ignore")
-register_sql_engine(
-    "qpdpandas", lambda engine: QPDPandasEngine(engine), on_dup="ignore"
-)
-register_sql_engine(
-    "qpd_pandas", lambda engine: QPDPandasEngine(engine), on_dup="ignore"
-)
+
+def register():
+    register_execution_engine(
+        "native", lambda conf: NativeExecutionEngine(conf), on_dup="ignore"
+    )
+    register_execution_engine(
+        "pandas", lambda conf: NativeExecutionEngine(conf), on_dup="ignore"
+    )
+    register_sql_engine("sqlite", lambda engine: SqliteEngine(engine), on_dup="ignore")
+    register_sql_engine(
+        "qpdpandas", lambda engine: QPDPandasEngine(engine), on_dup="ignore"
+    )
+    register_sql_engine(
+        "qpd_pandas", lambda engine: QPDPandasEngine(engine), on_dup="ignore"
+    )
+
+
+register()

fugue/execution/factory.py

@@ -352,6 +352,8 @@ def make_execution_engine(
             # SparkExecutionEngine + S2
             make_execution_engine((SparkExecutionEngine, "s"))
     """
+    import fugue._utils.register  # pylint: disable=W0611 # noqa: F401
+
     return _EXECUTION_ENGINE_FACTORY.make(engine, conf, **kwargs)
 
 
@@ -404,4 +406,6 @@ def make_sql_engine(
             # SqliteEngine(engine)
             make_sql_engine(SqliteEngine)
     """
+    import fugue._utils.register  # pylint: disable=W0611 # noqa: F401
+
     return _EXECUTION_ENGINE_FACTORY.make_sql_engine(engine, execution_engine, **kwargs)

fugue/extensions/_utils.py

@@ -27,6 +27,8 @@ def register(self, name: str, extension: Any, on_dup="overwrite") -> None:
         raise ValueError(on_dup)
 
     def get(self, obj: Any) -> Any:
+        import fugue._utils.register  # pylint: disable=W0611 # noqa: F401
+
         if isinstance(obj, str) and obj in self._dict:
             return self._dict[obj]
         return obj

fugue/workflow/utils.py

@@ -23,4 +23,6 @@ def is_acceptable_raw_df(df: Any) -> bool:
     :param df: input raw dataframe
     :return: whether this dataframe is convertible
     """
+    import fugue._utils.register  # pylint: disable=W0611 # noqa: F401
+
     return any(isinstance(df, t) for t in _VALID_RAW_DF_TYPES)

fugue_ibis/__init__.py

@@ -1,4 +1,8 @@
 # flake8: noqa
-from fugue_ibis.execution import pandas_backend
-from fugue_ibis.execution.ibis_engine import IbisEngine
-from fugue_ibis.extensions import run_ibis, as_ibis, as_fugue
+from fugue_ibis.execution.ibis_engine import IbisEngine, register_ibis_engine
+from fugue_ibis.execution.pandas_backend import _to_pandas_ibis_engine
+from fugue_ibis.extensions import as_fugue, as_ibis, run_ibis
+
+
+def register():
+    register_ibis_engine(1, _to_pandas_ibis_engine)

fugue_ibis/execution/pandas_backend.py

@@ -12,7 +12,7 @@
 from triad.utils.assertion import assert_or_throw
 import pandas as pd
 
-from fugue_ibis.execution.ibis_engine import IbisEngine, register_ibis_engine
+from fugue_ibis.execution.ibis_engine import IbisEngine
 from fugue_ibis._utils import to_schema, to_ibis_schema
 from ibis.backends.pandas import Backend
 
@@ -53,6 +53,3 @@ def table(self, name: str, schema: Any = None):
             if schema is None and name in self._schemas
             else schema,
         )
-
-
-register_ibis_engine(1, _to_pandas_ibis_engine)

fugue_notebook/env.py

@@ -5,12 +5,7 @@
 
 import fugue_sql
 import pandas as pd
-from fugue import (
-    ExecutionEngine,
-    NativeExecutionEngine,
-    make_execution_engine,
-    register_execution_engine,
-)
+from fugue import ExecutionEngine, make_execution_engine
 from fugue.dataframe import YieldedDataFrame
 from fugue.extensions._builtins.outputters import Show
 from fugue_sql.exceptions import FugueSQLSyntaxError
@@ -38,33 +33,6 @@ def get_pretty_print(self) -> Callable:
         """Fugue dataframe pretty print handler"""
         return _default_pretty_print
 
-    def register_execution_engines(self):
-        """Register execution engines with names. This will also try to register
-        spark and dask engines if the dependent packages are available and they
-        are not registered"""
-        register_execution_engine(
-            "native",
-            lambda conf, **kwargs: NativeExecutionEngine(conf=conf),
-            on_dup="ignore",
-        )
-
-        try:
-            import pyspark  # noqa: F401
-            import fugue_spark  # noqa: F401
-        except ImportError:
-            pass
-
-        try:
-            import dask.dataframe  # noqa: F401
-            import fugue_dask  # noqa: F401
-        except ImportError:
-            pass
-
-        try:
-            import fugue_duckdb  # noqa: F401
-        except ImportError:
-            pass
-
 
 @magics_class
 class _FugueSQLMagics(Magics):
@@ -151,5 +119,4 @@ def _setup_fugue_notebook(
         fsql_ignore_case=fsql_ignore_case,
     )
     ipython.register_magics(magics)
-    s.register_execution_engines()
     Show.set_hook(s.get_pretty_print())

fugue_version/__init__.py

@@ -1 +1 @@
-__version__ = "0.6.5"
+__version__ = "0.6.6"

setup.py

@@ -35,6 +35,7 @@ def get_version() -> str:
         "sqlalchemy",
         "pyarrow>=0.15.1",
         "pandas>=1.0.2",
+        "importlib-metadata; python_version < '3.8'",
     ],
     extras_require={
         "sql": ["antlr4-python3-runtime", "jinja2"],
@@ -71,4 +72,12 @@ def get_version() -> str:
     ],
     python_requires=">=3.6",
     package_data={"fugue_notebook": ["nbextension/*"]},
+    entry_points={
+        "fugue.plugins": [
+            "ibis = fugue_ibis:register",
+            "duckdb = fugue_duckdb:register",
+            "spark = fugue_spark:register",
+            "dask = fugue_dask:register",
+        ]
+    },
 )

tests/fugue_dask/test_importless.py

@@ -0,0 +1,24 @@
+from fugue import FugueWorkflow
+from fugue_sql import fsql
+
+
+def test_importless():
+    for engine in ["dask"]:
+        dag = FugueWorkflow()
+        dag.df([[0]], "a:int").show()
+
+        dag.run(engine)
+
+        fsql(
+            """
+        CREATE [[0],[1]] SCHEMA a:int
+        SELECT * WHERE a<1
+        PRINT
+        """
+        ).run(engine)
+
+        dag = FugueWorkflow()
+        idf = dag.df([[0], [1]], "a:int").as_ibis()
+        idf[idf.a < 1].as_fugue().show()
+
+        dag.run(engine)

tests/fugue_duckdb/test_importless.py

@@ -0,0 +1,30 @@
+from fugue import FugueWorkflow
+from fugue_sql import fsql
+
+
+def test_importless():
+    for engine in ["duck", "duckdb"]:
+        dag = FugueWorkflow()
+        dag.df([[0]], "a:int").show()
+
+        dag.run(engine)
+
+        fsql(
+            """
+        CREATE [[0],[1]] SCHEMA a:int
+        SELECT * WHERE a<1
+        PRINT
+        """
+        ).run(engine)
+
+        dag = FugueWorkflow()
+        idf = dag.df([[0], [1]], "a:int").as_ibis()
+        idf[idf.a < 1].as_fugue().show()
+
+        dag.run(engine)
+
+        dag = FugueWorkflow()
+        tdf = dag.df([[0], [1]], "a:int")
+        dag.select("SELECT * FROM ", tdf, " WHERE a<1", sql_engine=engine)
+
+        dag.run()

tests/fugue_ibis/test_importless.py

@@ -0,0 +1,10 @@
+from fugue import FugueWorkflow
+
+
+def test_importless():
+    for engine in [None]:
+        dag = FugueWorkflow()
+        idf = dag.df([[0], [1]], "a:int").as_ibis()
+        idf[idf.a < 1].as_fugue().show()
+
+        dag.run(engine)

tests/fugue_spark/test_importless.py

@@ -0,0 +1,26 @@
+from fugue import FugueWorkflow
+from fugue_sql import fsql
+from pyspark.sql import SparkSession
+
+
+def test_importless():
+    spark = SparkSession.builder.getOrCreate()
+    for engine in [spark, "spark"]:
+        dag = FugueWorkflow()
+        dag.df([[0]], "a:int").show()
+
+        dag.run(engine)
+
+        fsql(
+            """
+        CREATE [[0],[1]] SCHEMA a:int
+        SELECT * WHERE a<1
+        PRINT
+        """
+        ).run(engine)
+
+        dag = FugueWorkflow()
+        idf = dag.df([[0], [1]], "a:int").as_ibis()
+        idf[idf.a < 1].as_fugue().show()
+
+        dag.run(engine)