jenkins-bot merged this change.

View Change

Make sure entites are preloaded on a DataSite

This resulted in AttributeError.

Change-Id: Iab08428441c35d8b762a55e89e8b182f3a47245e
---
M pywikibot/pagegenerators.py
M tests/pagegenerators_tests.py
M tests/wikibase_tests.py
3 files changed, 30 insertions(+), 11 deletions(-)

diff --git a/pywikibot/pagegenerators.py b/pywikibot/pagegenerators.py
index 55a5c1b..aeddfe5 100644
--- a/pywikibot/pagegenerators.py
+++ b/pywikibot/pagegenerators.py
@@ -2081,6 +2081,7 @@
     Function basically is copied from above, but for Wikibase entities.
 
     @param generator: pages to iterate over
+    @type generator: Iterable
     @param groupsize: how many pages to preload at once
     @type groupsize: int
     """
@@ -2091,11 +2092,13 @@
         if len(sites[site]) >= groupsize:
             # if this site is at the groupsize, process it
             group = sites.pop(site)
-            for i in site.preload_entities(group, groupsize):
+            repo = site.data_repository()
+            for i in repo.preload_entities(group, groupsize):
                 yield i
     for site, pages in sites.items():
         # process any leftover sites that never reached the groupsize
-        for i in site.preload_entities(pages, groupsize):
+        repo = site.data_repository()
+        for i in repo.preload_entities(pages, groupsize):
             yield i
 
 
diff --git a/tests/pagegenerators_tests.py b/tests/pagegenerators_tests.py
index 43ca6ce..d8d5a0c 100755
--- a/tests/pagegenerators_tests.py
+++ b/tests/pagegenerators_tests.py
@@ -600,16 +600,16 @@
         self.assertTrue(pages_out[1].isTalkPage())
 
 
-class TestPreloadingItemGenerator(WikidataTestCase):
+class TestPreloadingEntityGenerator(WikidataTestCase):
 
     """Test preloading item generator."""
 
     def test_non_item_gen(self):
-        """Test TestPreloadingItemGenerator with ReferringPageGenerator."""
+        """Test TestPreloadingEntityGenerator with ReferringPageGenerator."""
         site = self.get_site()
         instance_of_page = pywikibot.Page(site, 'Property:P31')
         ref_gen = pagegenerators.ReferringPageGenerator(instance_of_page, total=5)
-        gen = pagegenerators.PreloadingItemGenerator(ref_gen)
+        gen = pagegenerators.PreloadingEntityGenerator(ref_gen)
         self.assertTrue(all(isinstance(item, pywikibot.ItemPage) for item in gen))
 
 
diff --git a/tests/wikibase_tests.py b/tests/wikibase_tests.py
index f6bb342..2fc7fee 100644
--- a/tests/wikibase_tests.py
+++ b/tests/wikibase_tests.py
@@ -1647,19 +1647,35 @@
         self.assertEqual(response, self.data_out)
 
 
-class TestPreloadingItemGenerator(TestCase):
+class TestPreloadingEntityGenerator(TestCase):
 
     """Test preloading item generator."""
 
-    family = 'wikidata'
-    code = 'wikidata'
+    sites = {
+        'wikidata': {
+            'family': 'wikidata',
+            'code': 'wikidata',
+        },
+        'enwiki': {
+            'family': 'wikipedia',
+            'code': 'en',
+        }
+    }
 
     def test_non_item_gen(self):
-        """Test TestPreloadingItemGenerator with ReferringPageGenerator."""
-        site = self.get_site()
+        """Test PreloadingEntityGenerator with ReferringPageGenerator."""
+        site = self.get_site('wikidata')
         instance_of_page = pywikibot.Page(site, 'Property:P31')
         ref_gen = pagegenerators.ReferringPageGenerator(instance_of_page, total=5)
-        gen = pagegenerators.PreloadingItemGenerator(ref_gen)
+        gen = pagegenerators.PreloadingEntityGenerator(ref_gen)
+        self.assertTrue(all(isinstance(item, ItemPage) for item in gen))
+
+    def test_foreign_page_item_gen(self):
+        """Test PreloadingEntityGenerator with connected pages."""
+        site = self.get_site('enwiki')
+        page_gen = [pywikibot.Page(site, 'Main Page'),
+                    pywikibot.Page(site, 'New York City')]
+        gen = pagegenerators.PreloadingEntityGenerator(page_gen)
         self.assertTrue(all(isinstance(item, ItemPage) for item in gen))
 
 

To view, visit change 421539. To unsubscribe, visit settings.