From 01fb44fd40b2aef50ac2d1de7a375c745698343d Mon Sep 17 00:00:00 2001 From: Cristian Date: Tue, 8 Sep 2020 09:36:06 -0500 Subject: [PATCH] refactor: Change archive_links check to focus on queryset, so it allows other iterables and not just lists --- archivebox/extractors/__init__.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/archivebox/extractors/__init__.py b/archivebox/extractors/__init__.py index d10d3ab1..7820c4a3 100644 --- a/archivebox/extractors/__init__.py +++ b/archivebox/extractors/__init__.py @@ -4,6 +4,7 @@ import os from typing import Optional, List, Iterable from datetime import datetime +from django.db.models import QuerySet from ..index.schema import Link from ..index import ( @@ -131,14 +132,14 @@ def archive_link(link: Link, overwrite: bool=False, methods: Optional[Iterable[s @enforce_types def archive_links(all_links: any, overwrite: bool=False, methods: Optional[Iterable[str]]=None, out_dir: Optional[str]=None) -> List[Link]: - if type(all_links) is list: - num_links: int = len(all_links) - get_link = lambda x: x - get_iter = lambda x: x - else: + if type(all_links) is QuerySet: num_links: int = all_links.count() get_link = lambda x: x.as_link() get_iter = lambda x: x.iterator() + else: + num_links: int = len(all_links) + get_link = lambda x: x + get_iter = lambda x: x if num_links == 0: return []