Skip to content

Commit

Permalink
Mypy fixes
Browse files Browse the repository at this point in the history
  • Loading branch information
vblagoje committed Oct 8, 2023
1 parent a5cdbc7 commit 57ab88b
Showing 1 changed file with 2 additions and 3 deletions.
5 changes: 2 additions & 3 deletions haystack/preview/components/routers/file_router.py
Original file line number Diff line number Diff line change
Expand Up @@ -55,21 +55,20 @@ def from_dict(cls, data: Dict[str, Any]) -> "FileTypeRouter":
"""
return default_from_dict(cls, data)

def run(self, sources: List[Union[str, Path, ByteStream]]) -> Dict[str, List[Path]]:
def run(self, sources: List[Union[str, Path, ByteStream]]) -> Dict[str, List[Union[ByteStream, Path]]]:
"""
Categorizes the provided data sources by their MIME types.
:param: sources: A list of file paths or byte streams to categorize.
:return: A dictionary where keys are MIME types and values are lists of data sources.
"""

mime_types = defaultdict(list)
for source in sources:
if isinstance(source, str):
source = Path(source)

if isinstance(source, (str, Path)):
if isinstance(source, Path):
mime_type = self.get_mime_type(source)
elif isinstance(source, ByteStream):
mime_type = source.metadata.get("content_type")
Expand Down

0 comments on commit 57ab88b

Please sign in to comment.