Replaced Union types with Optional

Author: alexthomas93

src/neo4j_graphrag/experimental/components/entity_relation_extractor.py

@@ -291,7 +291,7 @@ async def run(
         chunks: TextChunks,
         document_info: Optional[DocumentInfo] = None,
         lexical_graph_config: Optional[LexicalGraphConfig] = None,
-        schema: Union[GraphSchema, None] = None,
+        schema: Optional[GraphSchema] = None,
         examples: str = "",
         **kwargs: Any,
     ) -> Neo4jGraph:

src/neo4j_graphrag/experimental/pipeline/config/base.py

@@ -17,7 +17,7 @@
 from __future__ import annotations
 
 import logging
-from typing import Any, Union
+from typing import Any, Optional
 
 from pydantic import BaseModel, PrivateAttr
 
@@ -58,5 +58,5 @@ def resolve_params(self, params: dict[str, ParamConfig]) -> dict[str, Any]:
             for param_name, param in params.items()
         }
 
-    def parse(self, resolved_data: Union[dict[str, Any], None] = None) -> Any:
+    def parse(self, resolved_data: Optional[dict[str, Any]] = None) -> Any:
         raise NotImplementedError()

src/neo4j_graphrag/experimental/pipeline/config/object_config.py

@@ -64,7 +64,7 @@ class ObjectConfig(AbstractConfig, Generic[T]):
     and its constructor parameters.
     """
 
-    class_: Union[str, None] = Field(default=None, validate_default=True)
+    class_: Optional[str] = Field(default=None, validate_default=True)
     """Path to class to be instantiated."""
     params_: dict[str, ParamConfig] = {}
     """Initialization parameters."""
@@ -119,7 +119,7 @@ def _get_class(cls, class_path: str, optional_module: Optional[str] = None) -> t
             raise ValueError(f"Could not find {class_name} in {module_name}")
         return cast(type, klass)
 
-    def parse(self, resolved_data: Union[dict[str, Any], None] = None) -> T:
+    def parse(self, resolved_data: Optional[dict[str, Any]] = None) -> T:
         """Import `class_`, resolve `params_` and instantiate object."""
         self._global_data = resolved_data or {}
         logger.debug(f"OBJECT_CONFIG: parsing {self} using {resolved_data}")
@@ -153,7 +153,7 @@ def validate_class(cls, class_: Any) -> str:
         # not used
         return "not used"
 
-    def parse(self, resolved_data: Union[dict[str, Any], None] = None) -> neo4j.Driver:
+    def parse(self, resolved_data: Optional[dict[str, Any]] = None) -> neo4j.Driver:
         params = self.resolve_params(self.params_)
         # we know these params are there because of the required params validator
         uri = params.pop("uri")
@@ -176,7 +176,7 @@ class Neo4jDriverType(RootModel):  # type: ignore[type-arg]
 
     model_config = ConfigDict(arbitrary_types_allowed=True)
 
-    def parse(self, resolved_data: Union[dict[str, Any], None] = None) -> neo4j.Driver:
+    def parse(self, resolved_data: Optional[dict[str, Any]] = None) -> neo4j.Driver:
         if isinstance(self.root, neo4j.Driver):
             return self.root
         # self.root is a Neo4jDriverConfig object
@@ -203,7 +203,7 @@ class LLMType(RootModel):  # type: ignore[type-arg]
 
     model_config = ConfigDict(arbitrary_types_allowed=True)
 
-    def parse(self, resolved_data: Union[dict[str, Any], None] = None) -> LLMInterface:
+    def parse(self, resolved_data: Optional[dict[str, Any]] = None) -> LLMInterface:
         if isinstance(self.root, LLMInterface):
             return self.root
         return self.root.parse(resolved_data)
@@ -229,7 +229,7 @@ class EmbedderType(RootModel):  # type: ignore[type-arg]
 
     model_config = ConfigDict(arbitrary_types_allowed=True)
 
-    def parse(self, resolved_data: Union[dict[str, Any], None] = None) -> Embedder:
+    def parse(self, resolved_data: Optional[dict[str, Any]] = None) -> Embedder:
         if isinstance(self.root, Embedder):
             return self.root
         return self.root.parse(resolved_data)
@@ -257,7 +257,7 @@ class ComponentType(RootModel):  # type: ignore[type-arg]
 
     model_config = ConfigDict(arbitrary_types_allowed=True)
 
-    def parse(self, resolved_data: Union[dict[str, Any], None] = None) -> Component:
+    def parse(self, resolved_data: Optional[dict[str, Any]] = None) -> Component:
         if isinstance(self.root, Component):
             return self.root
         return self.root.parse(resolved_data)

src/neo4j_graphrag/experimental/pipeline/config/runner.py

@@ -24,6 +24,7 @@
 from typing import (
     Annotated,
     Any,
+    Optional,
     Union,
 )
 
@@ -71,7 +72,7 @@ class PipelineConfigWrapper(BaseModel):
     ] = Field(discriminator=Discriminator(_get_discriminator_value))
 
     def parse(
-        self, resolved_data: Union[dict[str, Any], None] = None
+        self, resolved_data: Optional[dict[str, Any]] = None
     ) -> PipelineDefinition:
         logger.debug("PIPELINE_CONFIG: start parsing config...")
         return self.config.parse(resolved_data)

src/neo4j_graphrag/experimental/pipeline/orchestrator.py

@@ -19,7 +19,7 @@
 import uuid
 import warnings
 from functools import partial
-from typing import TYPE_CHECKING, Any, AsyncGenerator, Union
+from typing import TYPE_CHECKING, Any, AsyncGenerator, Optional
 
 from neo4j_graphrag.experimental.pipeline.exceptions import (
     PipelineDefinitionError,
@@ -235,7 +235,7 @@ async def get_component_inputs(
         return component_inputs
 
     async def add_result_for_component(
-        self, name: str, result: Union[dict[str, Any], None], is_final: bool = False
+        self, name: str, result: Optional[dict[str, Any]], is_final: bool = False
     ) -> None:
         """This is where we save the results in the result store and, optionally,
         in the final result store.

src/neo4j_graphrag/experimental/pipeline/pipeline.py

@@ -19,7 +19,7 @@
 import warnings
 from collections import defaultdict
 from timeit import default_timer
-from typing import Any, AsyncGenerator, Optional, Union
+from typing import Any, AsyncGenerator, Optional
 
 from neo4j_graphrag.utils.logging import prettify
 
@@ -78,7 +78,7 @@ def __init__(self, name: str, component: Component):
 
     async def execute(
         self, context: RunContext, inputs: dict[str, Any]
-    ) -> Union[RunResult, None]:
+    ) -> Optional[RunResult]:
         """Execute the task
 
         Returns:

src/neo4j_graphrag/generation/graphrag.py

@@ -88,8 +88,8 @@ def search(
         message_history: Optional[Union[List[LLMMessage], MessageHistory]] = None,
         examples: str = "",
         retriever_config: Optional[dict[str, Any]] = None,
-        return_context: Union[bool, None] = None,
-        response_fallback: Union[str, None] = None,
+        return_context: Optional[bool] = None,
+        response_fallback: Optional[str] = None,
     ) -> RagResultModel:
         """
         .. warning::

src/neo4j_graphrag/generation/types.py

@@ -14,7 +14,7 @@
 #  limitations under the License.
 from __future__ import annotations
 
-from typing import Any, Union
+from typing import Any, Optional
 
 from pydantic import BaseModel, ConfigDict, field_validator
 
@@ -43,11 +43,11 @@ class RagSearchModel(BaseModel):
     examples: str = ""
     retriever_config: dict[str, Any] = {}
     return_context: bool = False
-    response_fallback: Union[str, None] = None
+    response_fallback: Optional[str] = None
 
 
 class RagResultModel(BaseModel):
     answer: str
-    retriever_result: Union[RetrieverResult, None] = None
+    retriever_result: Optional[RetrieverResult] = None
 
     model_config = ConfigDict(arbitrary_types_allowed=True)

tests/unit/experimental/components/test_graph_pruning.py

@@ -14,7 +14,7 @@
 #  limitations under the License.
 from __future__ import annotations
 
-from typing import Any, Union
+from typing import Any, Optional
 from unittest.mock import ANY, Mock, patch
 
 import pytest
@@ -369,7 +369,7 @@ def test_graph_pruning_validate_relationship(
     additional_relationship_types: bool,
     patterns: tuple[tuple[str, str, str], ...],
     additional_patterns: bool,
-    expected_relationship: Union[str, None],
+    expected_relationship: Optional[str],
     request: pytest.FixtureRequest,
 ) -> None:
     relationship_obj = request.getfixturevalue(relationship)

tests/unit/retrievers/test_base.py

@@ -15,7 +15,7 @@
 from __future__ import annotations  # Reminder: May be removed after Python 3.9 is EOL.
 
 import inspect
-from typing import Any, Union
+from typing import Any, Optional
 from unittest.mock import MagicMock, patch
 
 import pytest
@@ -41,7 +41,7 @@ def test_retriever_version_support(
     mock_get_version: MagicMock,
     driver: MagicMock,
     db_version: tuple[tuple[int, ...], bool],
-    expected_exception: Union[type[ValueError], None],
+    expected_exception: Optional[type[ValueError]],
 ) -> None:
     mock_get_version.return_value = db_version