diff --git a/hgext/largefiles/lfutil.py b/hgext/largefiles/lfutil.py --- a/hgext/largefiles/lfutil.py +++ b/hgext/largefiles/lfutil.py @@ -159,6 +159,8 @@ def findfile(repo, hash): class largefilesdirstate(dirstate.dirstate): + _large_file_dirstate = True + def __getitem__(self, key): return super(largefilesdirstate, self).__getitem__(unixpath(key)) @@ -204,7 +206,13 @@ def openlfdirstate(ui, repo, create=True """ Return a dirstate object that tracks largefiles: i.e. its root is the repo root, but it is saved in .hg/largefiles/dirstate. + + If a dirstate object already exists and is being used for a 'changing_*' + context, it will be returned. """ + sub_dirstate = getattr(repo.dirstate, '_sub_dirstate', None) + if sub_dirstate is not None: + return sub_dirstate vfs = repo.vfs lfstoredir = longname opener = vfsmod.vfs(vfs.join(lfstoredir)) diff --git a/hgext/largefiles/overrides.py b/hgext/largefiles/overrides.py --- a/hgext/largefiles/overrides.py +++ b/hgext/largefiles/overrides.py @@ -8,6 +8,7 @@ '''Overridden Mercurial commands and functions for the largefiles extension''' +import contextlib import copy import os @@ -21,6 +22,7 @@ from mercurial import ( archival, cmdutil, copies as copiesmod, + dirstate, error, exchange, extensions, @@ -311,6 +313,27 @@ def cmdutilremove( ) +@eh.wrapfunction(dirstate.dirstate, b'_changing') +@contextlib.contextmanager +def _changing(orig, self, repo, change_type): + pre = sub_dirstate = getattr(self, '_sub_dirstate', None) + try: + lfd = getattr(self, '_large_file_dirstate', False) + if sub_dirstate is None and not lfd: + sub_dirstate = lfutil.openlfdirstate(repo.ui, repo) + self._sub_dirstate = sub_dirstate + if not lfd: + assert self._sub_dirstate is not None + with orig(self, repo, change_type): + if sub_dirstate is None: + yield + else: + with sub_dirstate._changing(repo, change_type): + yield + finally: + self._sub_dirstate = pre + + @eh.wrapfunction(subrepo.hgsubrepo, b'status') def overridestatusfn(orig, repo, rev2, **opts): with lfstatus(repo._repo):