Skip to content

Commit

Permalink
Adjust typing in Mongo hook to prevent mypy errors (#42354)
Browse files Browse the repository at this point in the history
  • Loading branch information
jscheffl authored Sep 19, 2024
1 parent f278e62 commit 4da35b5
Showing 1 changed file with 4 additions and 4 deletions.
8 changes: 4 additions & 4 deletions airflow/providers/mongo/hooks/mongo.py
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,8 @@
if TYPE_CHECKING:
from types import TracebackType

from pymongo.collection import Collection as MongoCollection
from pymongo.command_cursor import CommandCursor
from typing_extensions import Literal

from airflow.models import Connection
Expand Down Expand Up @@ -218,9 +220,7 @@ def _create_uri(self) -> str:
path = f"/{self.connection.schema}"
return urlunsplit((scheme, netloc, path, "", ""))

def get_collection(
self, mongo_collection: str, mongo_db: str | None = None
) -> pymongo.collection.Collection:
def get_collection(self, mongo_collection: str, mongo_db: str | None = None) -> MongoCollection:
"""
Fetch a mongo collection object for querying.
Expand All @@ -233,7 +233,7 @@ def get_collection(

def aggregate(
self, mongo_collection: str, aggregate_query: list, mongo_db: str | None = None, **kwargs
) -> pymongo.command_cursor.CommandCursor:
) -> CommandCursor:
"""
Run an aggregation pipeline and returns the results.
Expand Down

0 comments on commit 4da35b5

Please sign in to comment.