Add entrypoint to setup.py in flytekit plugins (#1120)

Signed-off-by: Kevin Su <[email protected]>
flyteorg · Aug 4, 2022 · 9b3aa0e · 9b3aa0e
1 parent e407a9d
commit 9b3aa0e
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 6 deletions.
diff --git a/plugins/flytekit-polars/setup.py b/plugins/flytekit-polars/setup.py
@@ -35,4 +35,5 @@
         "Topic :: Software Development :: Libraries",
         "Topic :: Software Development :: Libraries :: Python Modules",
     ],
+    entry_points={"flytekit.plugins": [f"{PLUGIN_NAME}=flytekitplugins.{PLUGIN_NAME}"]},
 )
diff --git a/plugins/flytekit-polars/tests/test_polars_plugin_sd.py b/plugins/flytekit-polars/tests/test_polars_plugin_sd.py
@@ -1,10 +1,5 @@
-import flytekitplugins.polars  # noqa F401
 import polars as pl
-
-try:
-    from typing import Annotated
-except ImportError:
-    from typing_extensions import Annotated
+from typing_extensions import Annotated
 
 from flytekit import kwtypes, task, workflow
 from flytekit.types.structured.structured_dataset import PARQUET, StructuredDataset

diff --git a/plugins/flytekit-whylogs/setup.py b/plugins/flytekit-whylogs/setup.py
@@ -33,4 +33,5 @@
         "Topic :: Software Development :: Libraries",
         "Topic :: Software Development :: Libraries :: Python Modules",
     ],
+    entry_points={"flytekit.plugins": [f"{PLUGIN_NAME}=flytekitplugins.{PLUGIN_NAME}"]},
 )