Reduce requirements by deferring import until when it's actually needed (#8)

* Defer imports in huggingface.py until it's actually needed

* format

* fix

* make dependencies optional

* make dependencies required

Author: wengh

pyproject.toml

@@ -12,6 +12,7 @@ packages = [
 [tool.poetry.dependencies]
 python = "^3.9"
 datasets = "^3.2"
+huggingface_hub = "^0.27.1"
 
 [tool.poetry.group.dev.dependencies]
 pytest = "^8.0.0"

pyspark_huggingface/huggingface.py

@@ -1,10 +1,13 @@
-from typing import Optional
+from typing import TYPE_CHECKING, Optional
 
-from pyspark.sql.datasource import DataSource, DataSourceArrowWriter, DataSourceReader
-from pyspark.sql.types import StructType
+from pyspark.sql.datasource import DataSource
 
-from pyspark_huggingface.huggingface_sink import HuggingFaceSink
-from pyspark_huggingface.huggingface_source import HuggingFaceSource
+if TYPE_CHECKING:
+    from pyspark.sql.datasource import DataSourceWriter, DataSourceReader
+    from pyspark.sql.types import StructType
+
+    from pyspark_huggingface.huggingface_sink import HuggingFaceSink
+    from pyspark_huggingface.huggingface_source import HuggingFaceSource
 
 
 class HuggingFaceDatasets(DataSource):
@@ -24,15 +27,19 @@ class HuggingFaceDatasets(DataSource):
     def __init__(self, options: dict):
         super().__init__(options)
         self.options = options
-        self.source: Optional[HuggingFaceSource] = None
-        self.sink: Optional[HuggingFaceSink] = None
+        self.source: Optional["HuggingFaceSource"] = None
+        self.sink: Optional["HuggingFaceSink"] = None
+
+    def get_source(self) -> "HuggingFaceSource":
+        from pyspark_huggingface.huggingface_source import HuggingFaceSource
 
-    def get_source(self) -> HuggingFaceSource:
         if self.source is None:
             self.source = HuggingFaceSource(self.options.copy())
         return self.source
 
-    def get_sink(self):
+    def get_sink(self) -> "HuggingFaceSink":
+        from pyspark_huggingface.huggingface_sink import HuggingFaceSink
+
         if self.sink is None:
             self.sink = HuggingFaceSink(self.options.copy())
         return self.sink
@@ -44,8 +51,8 @@ def name(cls):
     def schema(self):
         return self.get_source().schema()
 
-    def reader(self, schema: StructType) -> "DataSourceReader":
+    def reader(self, schema: "StructType") -> "DataSourceReader":
         return self.get_source().reader(schema)
 
-    def writer(self, schema: StructType, overwrite: bool) -> "DataSourceArrowWriter":
+    def writer(self, schema: "StructType", overwrite: bool) -> "DataSourceWriter":
         return self.get_sink().writer(schema, overwrite)