From 30fd8c3464dd58895ad87d42dc584b2fee6ceb1c Mon Sep 17 00:00:00 2001 From: Emmanuelle Vargas-Gonzalez Date: Tue, 2 Feb 2021 00:08:11 -0500 Subject: [PATCH] compact calls --- stix2/datastore/taxii.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/stix2/datastore/taxii.py b/stix2/datastore/taxii.py index a0e2706..9ad6df9 100644 --- a/stix2/datastore/taxii.py +++ b/stix2/datastore/taxii.py @@ -297,12 +297,10 @@ class TAXIICollectionSource(DataSource): # query TAXII collection all_data = [] try: - if isinstance(self.collection, tcv21.Collection): - for envelope in tcv21.as_pages(self.collection.get_objects, per_request=self.items_per_page, **taxii_filters_dict): - all_data.extend(envelope.get("objects", [])) - else: - for bundle in tcv20.as_pages(self.collection.get_objects, per_request=self.items_per_page, **taxii_filters_dict): - all_data.extend(bundle.get("objects", [])) + paged_request = tcv21.as_pages if isinstance(self.collection, tcv21.Collection) else tcv20.as_pages + + for resource in paged_request(self.collection.get_objects, per_request=self.items_per_page, **taxii_filters_dict): + all_data.extend(resource.get("objects", [])) # deduplicate data (before filtering as reduces wasted filtering) all_data = deduplicate(all_data)