diff --git a/mercurial/setdiscovery.py b/mercurial/setdiscovery.py
--- a/mercurial/setdiscovery.py
+++ b/mercurial/setdiscovery.py
@@ -52,6 +52,7 @@ from .node import (
 )
 from . import (
     error,
+    policy,
     util,
 )
 
@@ -269,6 +270,10 @@ class partialdiscovery(object):
                 sample.update(takefrom[:more])
         return sample
 
+partialdiscovery = policy.importrust('discovery',
+                                     member='PartialDiscovery',
+                                     default=partialdiscovery)
+
 def findcommonheads(ui, local, remote,
                     initialsamplesize=100,
                     fullsamplesize=200,