narrow: only wrap dirstate functions once, instead of per-reposetup

chg will call reposetup multiple times, and we would end up double-wrapping (or
worse) the dirstate functions; this can cause issues like OSError 'No such file
or directory' during rebase operations, when we go to double-delete our
narrowspec backup file.

Differential Revision: https://phab.mercurial-scm.org/D3559

hgext/narrow/__init__.py

@@ -28,7 +28,6 b' from . import ('
     narrowchangegroup,
     narrowcommands,
     narrowcopies,
-    narrowdirstate,
     narrowpatch,
     narrowrepo,
     narrowrevlog,
@@ -72,10 +71,9 b' def reposetup(ui, repo):'
     if not repo.local():
         return
 
-    narrowrepo.wraprepo(repo)
     if changegroup.NARROW_REQUIREMENT in repo.requirements:
+        narrowrepo.wraprepo(repo)
         narrowcopies.setup(repo)
-        narrowdirstate.setup(repo)
         narrowpatch.setup(repo)
         narrowwirepeer.reposetup(repo)
 

hgext/narrow/narrowdirstate.py

@@ -9,74 +9,91 b' from __future__ import absolute_import'
 
 from mercurial.i18n import _
 from mercurial import (
-    dirstate,
     error,
-    extensions,
     match as matchmod,
     narrowspec,
     util as hgutil,
 )
 
-def setup(repo):
+def wrapdirstate(repo, dirstate):
     """Add narrow spec dirstate ignore, block changes outside narrow spec."""
 
-    def walk(orig, self, match, subrepos, unknown, ignored, full=True,
-             narrowonly=True):
-        if narrowonly:
-            # hack to not exclude explicitly-specified paths so that they can
-            # be warned later on e.g. dirstate.add()
-            em = matchmod.exact(match._root, match._cwd, match.files())
-            nm = matchmod.unionmatcher([repo.narrowmatch(), em])
-            match = matchmod.intersectmatchers(match, nm)
-        return orig(self, match, subrepos, unknown, ignored, full)
-
-    extensions.wrapfunction(dirstate.dirstate, 'walk', walk)
-
-    # Prevent adding files that are outside the sparse checkout
-    editfuncs = ['normal', 'add', 'normallookup', 'copy', 'remove', 'merge']
-    for func in editfuncs:
-        def _wrapper(orig, self, *args):
+    def _editfunc(fn):
+        def _wrapper(self, *args):
             dirstate = repo.dirstate
             narrowmatch = repo.narrowmatch()
             for f in args:
                 if f is not None and not narrowmatch(f) and f not in dirstate:
                     raise error.Abort(_("cannot track '%s' - it is outside " +
                         "the narrow clone") % f)
-            return orig(self, *args)
-        extensions.wrapfunction(dirstate.dirstate, func, _wrapper)
-
-    def filterrebuild(orig, self, parent, allfiles, changedfiles=None):
-        if changedfiles is None:
-            # Rebuilding entire dirstate, let's filter allfiles to match the
-            # narrowspec.
-            allfiles = [f for f in allfiles if repo.narrowmatch()(f)]
-        orig(self, parent, allfiles, changedfiles)
-
-    extensions.wrapfunction(dirstate.dirstate, 'rebuild', filterrebuild)
+            return fn(self, *args)
+        return _wrapper
 
     def _narrowbackupname(backupname):
         assert 'dirstate' in backupname
         return backupname.replace('dirstate', narrowspec.FILENAME)
 
-    def restorebackup(orig, self, tr, backupname):
-        self._opener.rename(_narrowbackupname(backupname), narrowspec.FILENAME,
-                            checkambig=True)
-        orig(self, tr, backupname)
+    class narrowdirstate(dirstate.__class__):
+        def walk(self, match, subrepos, unknown, ignored, full=True,
+                 narrowonly=True):
+            if narrowonly:
+                # hack to not exclude explicitly-specified paths so that they
+                # can be warned later on e.g. dirstate.add()
+                em = matchmod.exact(match._root, match._cwd, match.files())
+                nm = matchmod.unionmatcher([repo.narrowmatch(), em])
+                match = matchmod.intersectmatchers(match, nm)
+            return super(narrowdirstate, self).walk(match, subrepos, unknown,
+                                                    ignored, full)
 
-    extensions.wrapfunction(dirstate.dirstate, 'restorebackup', restorebackup)
+        # Prevent adding/editing/copying/deleting files that are outside the
+        # sparse checkout
+        @_editfunc
+        def normal(self, *args):
+            return super(narrowdirstate, self).normal(*args)
 
-    def savebackup(orig, self, tr, backupname):
-        orig(self, tr, backupname)
+        @_editfunc
+        def add(self, *args):
+            return super(narrowdirstate, self).add(*args)
+
+        @_editfunc
+        def normallookup(self, *args):
+            return super(narrowdirstate, self).normallookup(*args)
+
+        @_editfunc
+        def copy(self, *args):
+            return super(narrowdirstate, self).copy(*args)
 
-        narrowbackupname = _narrowbackupname(backupname)
-        self._opener.tryunlink(narrowbackupname)
-        hgutil.copyfile(self._opener.join(narrowspec.FILENAME),
-                        self._opener.join(narrowbackupname), hardlink=True)
+        @_editfunc
+        def remove(self, *args):
+            return super(narrowdirstate, self).remove(*args)
+
+        @_editfunc
+        def merge(self, *args):
+            return super(narrowdirstate, self).merge(*args)
+
+        def rebuild(self, parent, allfiles, changedfiles=None):
+            if changedfiles is None:
+                # Rebuilding entire dirstate, let's filter allfiles to match the
+                # narrowspec.
+                allfiles = [f for f in allfiles if repo.narrowmatch()(f)]
+            super(narrowdirstate, self).rebuild(parent, allfiles, changedfiles)
 
-    extensions.wrapfunction(dirstate.dirstate, 'savebackup', savebackup)
+        def restorebackup(self, tr, backupname):
+            self._opener.rename(_narrowbackupname(backupname),
+                                narrowspec.FILENAME, checkambig=True)
+            super(narrowdirstate, self).restorebackup(tr, backupname)
+
+        def savebackup(self, tr, backupname):
+            super(narrowdirstate, self).savebackup(tr, backupname)
 
-    def clearbackup(orig, self, tr, backupname):
-        orig(self, tr, backupname)
-        self._opener.unlink(_narrowbackupname(backupname))
+            narrowbackupname = _narrowbackupname(backupname)
+            self._opener.tryunlink(narrowbackupname)
+            hgutil.copyfile(self._opener.join(narrowspec.FILENAME),
+                            self._opener.join(narrowbackupname), hardlink=True)
 
-    extensions.wrapfunction(dirstate.dirstate, 'clearbackup', clearbackup)
+        def clearbackup(self, tr, backupname):
+            super(narrowdirstate, self).clearbackup(tr, backupname)
+            self._opener.unlink(_narrowbackupname(backupname))
+
+    dirstate.__class__ = narrowdirstate
+    return dirstate

hgext/narrow/narrowrepo.py

@@ -15,6 +15,7 b' from mercurial import ('
 )
 
 from . import (
+    narrowdirstate,
     narrowrevlog,
 )
 
@@ -62,4 +63,8 b' def wraprepo(repo):'
             return scmutil.status(modified, added, removed, deleted, unknown,
                                   ignored, clean)
 
+        def _makedirstate(self):
+            dirstate = super(narrowrepository, self)._makedirstate()
+            return narrowdirstate.wrapdirstate(self, dirstate)
+
     repo.__class__ = narrowrepository

mercurial/localrepo.py

@@ -778,6 +778,9 b' class localrepository(object):'
 
     @repofilecache('dirstate')
     def dirstate(self):
+        return self._makedirstate()
+
+    def _makedirstate(self):
         sparsematchfn = lambda: sparse.matcher(self)
 
         return dirstate.dirstate(self.vfs, self.ui, self.root,

tests/test-narrow-expanddirstate.t

@@ -72,29 +72,31 b' have this method available in narrowhg p'
   >   for f in repo[b'.'].manifest().walk(added):
   >     repo.dirstate.normallookup(f)
   > 
-  > def makeds(ui, repo):
-  >   def wrapds(orig, self):
-  >     ds = orig(self)
-  >     class expandingdirstate(ds.__class__):
-  >       @hgutil.propertycache
-  >       def _map(self):
-  >         ret = super(expandingdirstate, self)._map
-  >         with repo.wlock(), repo.lock(), repo.transaction(
-  >             b'expandnarrowspec'):
-  >           expandnarrowspec(ui, repo,
-  >                            encoding.environ.get(b'DIRSTATEINCLUDES'))
-  >         return ret
-  >     ds.__class__ = expandingdirstate
-  >     return ds
-  >   return wrapds
+  > def wrapds(ui, repo, ds):
+  >   class expandingdirstate(ds.__class__):
+  >     @hgutil.propertycache
+  >     def _map(self):
+  >       ret = super(expandingdirstate, self)._map
+  >       with repo.wlock(), repo.lock(), repo.transaction(
+  >           b'expandnarrowspec'):
+  >         expandnarrowspec(ui, repo,
+  >                          encoding.environ.get(b'DIRSTATEINCLUDES'))
+  >       return ret
+  >   ds.__class__ = expandingdirstate
+  >   return ds
   > 
   > def reposetup(ui, repo):
-  >   extensions.wrapfilecache(localrepo.localrepository, b'dirstate',
-  >                            makeds(ui, repo))
-  >   def overridepatch(orig, *args, **kwargs):
+  >   class expandingrepo(repo.__class__):
+  >     def _makedirstate(self):
+  >       dirstate = super(expandingrepo, self)._makedirstate()
+  >       return wrapds(ui, repo, dirstate)
+  >   repo.__class__ = expandingrepo
+  > 
+  > def extsetup(unused_ui):
+  >   def overridepatch(orig, ui, repo, *args, **kwargs):
   >     with repo.wlock():
   >       expandnarrowspec(ui, repo, encoding.environ.get(b'PATCHINCLUDES'))
-  >       return orig(*args, **kwargs)
+  >       return orig(ui, repo, *args, **kwargs)
   > 
   >   extensions.wrapfunction(patch, b'patch', overridepatch)
   > EOF