Fix mypy issues

Author: ali-sedaghatbaf

examples/customize/build_graph/components/loaders/custom_loader.py

@@ -1,7 +1,9 @@
 """Create a custom data loader to transform content into text."""
 
 from pathlib import Path
-from typing import Dict, Optional
+from typing import Dict, Optional, Union
+
+from fsspec import AbstractFileSystem
 
 from neo4j_graphrag.experimental.components.data_loader import DataLoader
 from neo4j_graphrag.experimental.components.types import DocumentInfo, LoadedDocument
@@ -10,10 +12,12 @@
 class MyLoader(DataLoader):
     async def run(
         self,
-        filepath: Path,
+        filepath: Union[str, Path],
         metadata: Optional[Dict[str, str]] = None,
+        fs: Optional[Union[AbstractFileSystem, str]] = None,
     ) -> LoadedDocument:
-        # Implement logic here
+        # Implement logic here; use ``fs`` when reading from non-local storage.
+        _ = fs
         return LoadedDocument(
             text="<extracted text>",
             document_info=DocumentInfo(

tests/unit/experimental/pipeline/config/template_pipeline/test_simple_kg_builder.py

@@ -33,7 +33,10 @@
 from neo4j_graphrag.experimental.components.text_splitters.fixed_size_splitter import (
     FixedSizeSplitter,
 )
-from neo4j_graphrag.experimental.pipeline.config.object_config import ComponentConfig
+from neo4j_graphrag.experimental.pipeline.config.object_config import (
+    ComponentConfig,
+    ComponentType,
+)
 from neo4j_graphrag.experimental.pipeline.config.template_pipeline import (
     SimpleKGPipelineConfig,
 )
@@ -70,7 +73,7 @@ def test_simple_kg_pipeline_config_pdf_loader_deprecated_maps_to_file_loader() -
     with pytest.warns(DeprecationWarning, match="pdf_loader"):
         config = SimpleKGPipelineConfig(
             from_file=True,
-            pdf_loader=my_loader,
+            pdf_loader=ComponentType(my_loader),
         )
     assert config._get_file_loader() == my_loader
 
@@ -79,24 +82,28 @@ def test_simple_kg_pipeline_config_pdf_loader_and_file_loader_conflict() -> None
     with pytest.raises(ValueError, match="pdf_loader"):
         SimpleKGPipelineConfig(
             from_file=True,
-            file_loader=PdfLoader(),  # type: ignore[call-arg]
-            pdf_loader=PdfLoader(),
+            file_loader=ComponentType(PdfLoader()),
+            pdf_loader=ComponentType(PdfLoader()),
         )
 
 
 def test_simple_kg_pipeline_config_file_loader_from_file_is_true_class_overwrite() -> (
     None
 ):
     my_file_loader = PdfLoader()
-    config = SimpleKGPipelineConfig(from_file=True, file_loader=my_file_loader)  # type: ignore
+    config = SimpleKGPipelineConfig(
+        from_file=True, file_loader=ComponentType(my_file_loader)
+    )
     assert config._get_file_loader() == my_file_loader
 
 
 def test_simple_kg_pipeline_config_file_loader_class_overwrite_but_from_file_is_false() -> (
     None
 ):
     my_file_loader = PdfLoader()
-    config = SimpleKGPipelineConfig(from_file=False, file_loader=my_file_loader)  # type: ignore
+    config = SimpleKGPipelineConfig(
+        from_file=False, file_loader=ComponentType(my_file_loader)
+    )
     assert config._get_file_loader() is None
 
 
@@ -111,7 +118,7 @@ def test_simple_kg_pipeline_config_file_loader_from_file_is_true_class_overwrite
     mock_component_parse.return_value = my_file_loader
     config = SimpleKGPipelineConfig(
         from_file=True,
-        file_loader=my_file_loader_config,  # type: ignore
+        file_loader=ComponentType(my_file_loader_config),
     )
     assert config._get_file_loader() == my_file_loader