--- a/hgext/sparse.py	Sun Jul 09 15:11:19 2017 +0200
+++ b/hgext/sparse.py	Fri Jul 07 11:51:10 2017 -0700
@@ -79,7 +79,6 @@
 from mercurial import (
     cmdutil,
     commands,
-    context,
     dirstate,
     error,
     extensions,
@@ -100,9 +99,6 @@
 cmdtable = {}
 command = registrar.command(cmdtable)
 
-def uisetup(ui):
-    _setupcommit(ui)
-
 def extsetup(ui):
     sparse.enabled = True
 
@@ -134,27 +130,6 @@
         raise AttributeError(_("type '%s' has no property '%s'") % (origcls,
                              propname))
 
-def _setupcommit(ui):
-    def _refreshoncommit(orig, self, node):
-        """Refresh the checkout when commits touch .hgsparse
-        """
-        orig(self, node)
-        repo = self._repo
-
-        ctx = repo[node]
-        profiles = sparse.patternsforrev(repo, ctx.rev())[2]
-
-        # profiles will only have data if sparse is enabled.
-        if set(profiles) & set(ctx.files()):
-            origstatus = repo.status()
-            origsparsematch = sparse.matcher(repo)
-            sparse.refreshwdir(repo, origstatus, origsparsematch, force=True)
-
-        sparse.prunetemporaryincludes(repo)
-
-    extensions.wrapfunction(context.committablectx, 'markcommitted',
-        _refreshoncommit)
-
 def _setuplog(ui):
     entry = commands.table['^log|history']
     entry[1].append(('', 'sparse', None,

--- a/mercurial/context.py	Sun Jul 09 15:11:19 2017 +0200
+++ b/mercurial/context.py	Fri Jul 07 11:51:10 2017 -0700
@@ -38,6 +38,7 @@
     repoview,
     revlog,
     scmutil,
+    sparse,
     subrepo,
     util,
 )
@@ -1803,6 +1804,11 @@
             match.bad = bad
         return match
 
+    def markcommitted(self, node):
+        super(workingctx, self).markcommitted(node)
+
+        sparse.aftercommit(self._repo, node)
+
 class committablefilectx(basefilectx):
     """A committablefilectx provides common functionality for a file context
     that wants the ability to commit, e.g. workingfilectx or memfilectx."""

--- a/mercurial/sparse.py	Sun Jul 09 15:11:19 2017 +0200
+++ b/mercurial/sparse.py	Fri Jul 07 11:51:10 2017 -0700
@@ -478,3 +478,19 @@
         dirstate.normallookup(file)
 
     return added, dropped, lookup
+
+def aftercommit(repo, node):
+    """Perform actions after a working directory commit."""
+    # This function is called unconditionally, even if sparse isn't
+    # enabled.
+    ctx = repo[node]
+
+    profiles = patternsforrev(repo, ctx.rev())[2]
+
+    # profiles will only have data if sparse is enabled.
+    if set(profiles) & set(ctx.files()):
+        origstatus = repo.status()
+        origsparsematch = matcher(repo)
+        refreshwdir(repo, origstatus, origsparsematch, force=True)
+
+    prunetemporaryincludes(repo)