From ff1f985a2a736f788c2551ec0a33281d0ffccd11 Mon Sep 17 00:00:00 2001 From: Christophe Bornet Date: Thu, 15 Feb 2024 15:45:22 +0100 Subject: [PATCH] community: Fix some mypy types in cassandra doc loader (#17570) Thank you! --- .../langchain_community/document_loaders/cassandra.py | 10 +++++----- .../document_loaders/test_cassandra.py | 4 ++-- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/libs/community/langchain_community/document_loaders/cassandra.py b/libs/community/langchain_community/document_loaders/cassandra.py index a3b7732c131..3cef56a1cbc 100644 --- a/libs/community/langchain_community/document_loaders/cassandra.py +++ b/libs/community/langchain_community/document_loaders/cassandra.py @@ -29,18 +29,18 @@ class CassandraLoader(BaseLoader): table: Optional[str] = None, session: Optional[Session] = None, keyspace: Optional[str] = None, - query: Optional[Union[str, Statement]] = None, + query: Union[str, Statement, None] = None, page_content_mapper: Callable[[Any], str] = str, metadata_mapper: Callable[[Any], dict] = lambda _: {}, *, - query_parameters: Union[dict, Sequence] = None, # type: ignore[assignment] + query_parameters: Union[dict, Sequence, None] = None, query_timeout: Optional[float] = _NOT_SET, # type: ignore[assignment] query_trace: bool = False, - query_custom_payload: dict = None, # type: ignore[assignment] + query_custom_payload: Optional[dict] = None, query_execution_profile: Any = _NOT_SET, query_paging_state: Any = None, - query_host: Host = None, - query_execute_as: str = None, # type: ignore[assignment] + query_host: Optional[Host] = None, + query_execute_as: Optional[str] = None, ) -> None: """ Document Loader for Apache Cassandra. diff --git a/libs/community/tests/integration_tests/document_loaders/test_cassandra.py b/libs/community/tests/integration_tests/document_loaders/test_cassandra.py index 5562188eced..a93a6abba68 100644 --- a/libs/community/tests/integration_tests/document_loaders/test_cassandra.py +++ b/libs/community/tests/integration_tests/document_loaders/test_cassandra.py @@ -2,7 +2,7 @@ Test of Cassandra document loader class `CassandraLoader` """ import os -from typing import Any +from typing import Any, Iterator import pytest from langchain_core.documents import Document @@ -14,7 +14,7 @@ CASSANDRA_TABLE = "docloader_test_table" @pytest.fixture(autouse=True, scope="session") -def keyspace() -> str: # type: ignore[misc] +def keyspace() -> Iterator[str]: import cassio from cassandra.cluster import Cluster from cassio.config import check_resolve_session, resolve_keyspace