jenkins-bot has submitted this change and it was merged.
Change subject: Make Site objects pickle-able ......................................................................
Make Site objects pickle-able
Since we can't pickle Lock objects, those are removed upon pickling, and restored when the object is unpacked.
Bug: 62094 Change-Id: I3298fa212ad99e1674067b42a1c081f8bd0f779e --- M pywikibot/site.py M tests/site_tests.py 2 files changed, 44 insertions(+), 1 deletion(-)
Approvals: Legoktm: Looks good to me, approved jenkins-bot: Verified
diff --git a/pywikibot/site.py b/pywikibot/site.py index 385ed16..5c28490 100644 --- a/pywikibot/site.py +++ b/pywikibot/site.py @@ -84,6 +84,8 @@ def __repr__(self): return 'LoginStatus(%s)' % (LoginStatus.name(self.state))
+_families = {} +
def Family(fam=None, fatal=True): """Import the named family. @@ -98,6 +100,9 @@ """ if fam is None: fam = config.family + if fam in _families: + return _families[fam] + try: myfamily = imp.load_source(fam, config.family_files[fam]) except (ImportError, KeyError): @@ -109,7 +114,8 @@ sys.exit(1) else: raise Error("Family %s does not exist" % fam) - return myfamily.Family() + _families[fam] = myfamily.Family() + return _families[fam]
class BaseSite(object): @@ -212,6 +218,19 @@ if self.family == other.family: return cmp(self.code, other.code) return cmp(self.family.name, other.family.name) + + def __getstate__(self): + """ Remove Lock based classes before pickling """ + new = self.__dict__.copy() + del new['_pagemutex'] + if '_throttle' in new: + del new['_throttle'] + return new + + def __setstate__(self, attrs): + """ Restore things removed in __getstate__ """ + self.__dict__.update(attrs) + self._pagemutex = threading.Lock()
def user(self): """Return the currently-logged in bot user, or None.""" @@ -676,6 +695,19 @@ return APISite(site['code'], site['code']) raise ValueError("Cannot parse a site out of %s." % dbname)
+ def __getstate__(self): + """ Remove Lock based classes before pickling """ + new = super(APISite, self).__getstate__() + del new['sitelock'] + del new['_msgcache'] + return new + + def __setstate__(self, attrs): + """ Restore things removed in __getstate__ """ + super(APISite, self).__setstate__(attrs) + self.sitelock = threading.Lock() + self._msgcache = {} + def _generator(self, gen_class, type_arg=None, namespaces=None, step=None, total=None, **args): """Convenience method that returns an API generator. diff --git a/tests/site_tests.py b/tests/site_tests.py index cca2db0..5e2cdc2 100644 --- a/tests/site_tests.py +++ b/tests/site_tests.py @@ -997,6 +997,17 @@ mysite.loadrevisions(mypage, step=5, total=12) self.assertEqual(len(mypage._revisions), 12)
+ def testPickleAbility(self): + import pickle + #fam = pywikibot.site.Family('wikipedia') + #site = pywikibot.Site('en', fam) + #pywikibot.output(site.family.__dict__) + #pickle.dumps(site) + site = pywikibot.Site('en', 'wikipedia') + #pywikibot.output(site.family.__dict__) + pickle.dumps(site) + self.assertTrue(True) # No exception thrown! +
class TestSiteLoadRevisions(PywikibotTestCase): """Test cases for Site.loadrevision() method."""