fastannotate: stop using the `pycompat.open()` shim

hgext/fastannotate/protocol.py

@@ -1,263 +1,262
 # Copyright 2016-present Facebook. All Rights Reserved.
 #
 # protocol: logic for a server providing fastannotate support
 #
 # This software may be used and distributed according to the terms of the
 # GNU General Public License version 2 or any later version.
 from __future__ import annotations
 
 import contextlib
 import os
 
 from mercurial.i18n import _
-from mercurial.pycompat import open
 from mercurial import (
     error,
     extensions,
     hg,
     util,
     wireprotov1peer,
     wireprotov1server,
 )
 from mercurial.utils import (
     urlutil,
 )
 from . import context
 
 # common
 
 
 def _getmaster(ui):
     """get the mainbranch, and enforce it is set"""
     master = ui.config(b'fastannotate', b'mainbranch')
     if not master:
         raise error.Abort(
             _(
                 b'fastannotate.mainbranch is required '
                 b'for both the client and the server'
             )
         )
     return master
 
 
 # server-side
 
 
 def _capabilities(orig, repo, proto):
     result = orig(repo, proto)
     result.append(b'getannotate')
     return result
 
 
 def _getannotate(repo, proto, path, lastnode):
     # output:
     #   FILE := vfspath + '\0' + str(size) + '\0' + content
     #   OUTPUT := '' | FILE + OUTPUT
     result = b''
     buildondemand = repo.ui.configbool(
         b'fastannotate', b'serverbuildondemand', True
     )
     with context.annotatecontext(repo, path) as actx:
         if buildondemand:
             # update before responding to the client
             master = _getmaster(repo.ui)
             try:
                 if not actx.isuptodate(master):
                     actx.annotate(master, master)
             except Exception:
                 # non-fast-forward move or corrupted. rebuild automically.
                 actx.rebuild()
                 try:
                     actx.annotate(master, master)
                 except Exception:
                     actx.rebuild()  # delete files
             finally:
                 # although the "with" context will also do a close/flush, we
                 # need to do it early so we can send the correct respond to
                 # client.
                 actx.close()
         # send back the full content of revmap and linelog, in the future we
         # may want to do some rsync-like fancy updating.
         # the lastnode check is not necessary if the client and the server
         # agree where the main branch is.
         if actx.lastnode != lastnode:
             for p in [actx.revmappath, actx.linelogpath]:
                 if not os.path.exists(p):
                     continue
-                with open(p, b'rb') as f:
+                with open(p, 'rb') as f:
                     content = f.read()
                 vfsbaselen = len(repo.vfs.base + b'/')
                 relpath = p[vfsbaselen:]
                 result += b'%s\0%d\0%s' % (relpath, len(content), content)
     return result
 
 
 def _registerwireprotocommand():
     if b'getannotate' in wireprotov1server.commands:
         return
     wireprotov1server.wireprotocommand(b'getannotate', b'path lastnode')(
         _getannotate
     )
 
 
 def serveruisetup(ui):
     _registerwireprotocommand()
     extensions.wrapfunction(wireprotov1server, '_capabilities', _capabilities)
 
 
 # client-side
 
 
 def _parseresponse(payload):
     result = {}
     i = 0
     l = len(payload) - 1
     state = 0  # 0: vfspath, 1: size
     vfspath = size = b''
     while i < l:
         ch = payload[i : i + 1]
         if ch == b'\0':
             if state == 1:
                 result[vfspath] = payload[i + 1 : i + 1 + int(size)]
                 i += int(size)
                 state = 0
                 vfspath = size = b''
             elif state == 0:
                 state = 1
         else:
             if state == 1:
                 size += ch
             elif state == 0:
                 vfspath += ch
         i += 1
     return result
 
 
 def peersetup(ui, peer):
     class fastannotatepeer(peer.__class__):
         @wireprotov1peer.batchable
         def getannotate(self, path, lastnode=None):
             if not self.capable(b'getannotate'):
                 ui.warn(_(b'remote peer cannot provide annotate cache\n'))
                 return None, None
             else:
                 args = {b'path': path, b'lastnode': lastnode or b''}
                 return args, _parseresponse
 
     peer.__class__ = fastannotatepeer
 
 
 @contextlib.contextmanager
 def annotatepeer(repo):
     ui = repo.ui
 
     remotedest = ui.config(b'fastannotate', b'remotepath', b'default')
     remotepath = urlutil.get_unique_pull_path_obj(
         b'fastannotate',
         ui,
         remotedest,
     )
     peer = hg.peer(ui, {}, remotepath)
 
     try:
         yield peer
     finally:
         peer.close()
 
 
 def clientfetch(repo, paths, lastnodemap=None, peer=None):
     """download annotate cache from the server for paths"""
     if not paths:
         return
 
     if peer is None:
         with annotatepeer(repo) as peer:
             return clientfetch(repo, paths, lastnodemap, peer)
 
     if lastnodemap is None:
         lastnodemap = {}
 
     ui = repo.ui
     results = []
     with peer.commandexecutor() as batcher:
         ui.debug(b'fastannotate: requesting %d files\n' % len(paths))
         for p in paths:
             results.append(
                 batcher.callcommand(
                     b'getannotate',
                     {b'path': p, b'lastnode': lastnodemap.get(p)},
                 )
             )
 
         for result in results:
             r = result.result()
             # TODO: pconvert these paths on the server?
             r = {util.pconvert(p): v for p, v in r.items()}
             for path in sorted(r):
                 # ignore malicious paths
                 if not path.startswith(b'fastannotate/') or b'/../' in (
                     path + b'/'
                 ):
                     ui.debug(
                         b'fastannotate: ignored malicious path %s\n' % path
                     )
                     continue
                 content = r[path]
                 if ui.debugflag:
                     ui.debug(
                         b'fastannotate: writing %d bytes to %s\n'
                         % (len(content), path)
                     )
                 repo.vfs.makedirs(os.path.dirname(path))
                 with repo.vfs(path, b'wb') as f:
                     f.write(content)
 
 
 def _filterfetchpaths(repo, paths):
     """return a subset of paths whose history is long and need to fetch linelog
     from the server. works with remotefilelog and non-remotefilelog repos.
     """
     threshold = repo.ui.configint(b'fastannotate', b'clientfetchthreshold', 10)
     if threshold <= 0:
         return paths
 
     result = []
     for path in paths:
         try:
             if len(repo.file(path)) >= threshold:
                 result.append(path)
         except Exception:  # file not found etc.
             result.append(path)
 
     return result
 
 
 def localreposetup(ui, repo):
     class fastannotaterepo(repo.__class__):
         def prefetchfastannotate(self, paths, peer=None):
             master = _getmaster(self.ui)
             needupdatepaths = []
             lastnodemap = {}
             try:
                 for path in _filterfetchpaths(self, paths):
                     with context.annotatecontext(self, path) as actx:
                         if not actx.isuptodate(master, strict=False):
                             needupdatepaths.append(path)
                             lastnodemap[path] = actx.lastnode
                 if needupdatepaths:
                     clientfetch(self, needupdatepaths, lastnodemap, peer)
             except Exception as ex:
                 # could be directory not writable or so, not fatal
                 self.ui.debug(b'fastannotate: prefetch failed: %r\n' % ex)
 
     repo.__class__ = fastannotaterepo
 
 
 def clientreposetup(ui, repo):
     _registerwireprotocommand()
     if repo.local():
         localreposetup(ui, repo)
     # TODO: this mutates global state, but only if at least one repo
     # has the extension enabled. This is probably bad for hgweb.
     if peersetup not in hg.wirepeersetupfuncs:
         hg.wirepeersetupfuncs.append(peersetup)

hgext/fastannotate/revmap.py

@@ -1,260 +1,259
 # Copyright 2016-present Facebook. All Rights Reserved.
 #
 # revmap: trivial hg hash - linelog rev bidirectional map
 #
 # This software may be used and distributed according to the terms of the
 # GNU General Public License version 2 or any later version.
 
 from __future__ import annotations
 
 import bisect
 import io
 import os
 import struct
 
 from mercurial.node import hex
-from mercurial.pycompat import open
 from mercurial import (
     error as hgerror,
 )
 from . import error
 
 # the revmap file format is straightforward:
 #
 #    8 bytes: header
 #    1 byte : flag for linelog revision 1
 #    ? bytes: (optional) '\0'-terminated path string
 #             only exists if (flag & renameflag) != 0
 #   20 bytes: hg hash for linelog revision 1
 #    1 byte : flag for linelog revision 2
 #    ? bytes: (optional) '\0'-terminated path string
 #   20 bytes: hg hash for linelog revision 2
 #   ....
 #
 # the implementation is kinda stupid: __init__ loads the whole revmap.
 # no laziness. benchmark shows loading 10000 revisions is about 0.015
 # seconds, which looks enough for our use-case. if this implementation
 # becomes a bottleneck, we can change it to lazily read the file
 # from the end.
 
 # whether the changeset is in the side branch. i.e. not in the linear main
 # branch but only got referenced by lines in merge changesets.
 sidebranchflag = 1
 
 # whether the changeset changes the file path (ie. is a rename)
 renameflag = 2
 
 # len(mercurial.node.nullid)
 _hshlen = 20
 
 
 class revmap:
     """trivial hg bin hash - linelog rev bidirectional map
 
     also stores a flag (uint8) for each revision, and track renames.
     """
 
     HEADER = b'REVMAP1\0'
 
     def __init__(self, path=None):
         """create or load the revmap, optionally associate to a file
 
         if path is None, the revmap is entirely in-memory. the caller is
         responsible for locking. concurrent writes to a same file is unsafe.
         the caller needs to make sure one file is associated to at most one
         revmap object at a time."""
         self.path = path
         self._rev2hsh = [None]
         self._rev2flag = [None]
         self._hsh2rev = {}
         # since rename does not happen frequently, do not store path for every
         # revision. self._renamerevs can be used for bisecting.
         self._renamerevs = [0]
         self._renamepaths = [b'']
         self._lastmaxrev = -1
         if path:
             if os.path.exists(path):
                 self._load()
             else:
                 # write the header so "append" can do incremental updates
                 self.flush()
 
     def copyfrom(self, rhs):
         """copy the map data from another revmap. do not affect self.path"""
         self._rev2hsh = rhs._rev2hsh[:]
         self._rev2flag = rhs._rev2flag[:]
         self._hsh2rev = rhs._hsh2rev.copy()
         self._renamerevs = rhs._renamerevs[:]
         self._renamepaths = rhs._renamepaths[:]
         self._lastmaxrev = -1
 
     @property
     def maxrev(self):
         """return max linelog revision number"""
         return len(self._rev2hsh) - 1
 
     def append(self, hsh, sidebranch=False, path=None, flush=False):
         """add a binary hg hash and return the mapped linelog revision.
         if flush is True, incrementally update the file.
         """
         if hsh in self._hsh2rev:
             raise error.CorruptedFileError(
                 b'%r is in revmap already' % hex(hsh)
             )
         if len(hsh) != _hshlen:
             raise hgerror.ProgrammingError(
                 b'hsh must be %d-char long' % _hshlen
             )
         idx = len(self._rev2hsh)
         flag = 0
         if sidebranch:
             flag |= sidebranchflag
         if path is not None and path != self._renamepaths[-1]:
             flag |= renameflag
             self._renamerevs.append(idx)
             self._renamepaths.append(path)
         self._rev2hsh.append(hsh)
         self._rev2flag.append(flag)
         self._hsh2rev[hsh] = idx
         if flush:
             self.flush()
         return idx
 
     def rev2hsh(self, rev):
         """convert linelog revision to hg hash. return None if not found."""
         if rev > self.maxrev or rev < 0:
             return None
         return self._rev2hsh[rev]
 
     def rev2flag(self, rev):
         """get the flag (uint8) for a given linelog revision.
         return None if revision does not exist.
         """
         if rev > self.maxrev or rev < 0:
             return None
         return self._rev2flag[rev]
 
     def rev2path(self, rev):
         """get the path for a given linelog revision.
         return None if revision does not exist.
         """
         if rev > self.maxrev or rev < 0:
             return None
         idx = bisect.bisect_right(self._renamerevs, rev) - 1
         return self._renamepaths[idx]
 
     def hsh2rev(self, hsh):
         """convert hg hash to linelog revision. return None if not found."""
         return self._hsh2rev.get(hsh)
 
     def clear(self, flush=False):
         """make the map empty. if flush is True, write to disk"""
         # rev 0 is reserved, real rev starts from 1
         self._rev2hsh = [None]
         self._rev2flag = [None]
         self._hsh2rev = {}
         self._rev2path = [b'']
         self._lastmaxrev = -1
         if flush:
             self.flush()
 
     def flush(self):
         """write the state down to the file"""
         if not self.path:
             return
         if self._lastmaxrev == -1:  # write the entire file
-            with open(self.path, b'wb') as f:
+            with open(self.path, 'wb') as f:
                 f.write(self.HEADER)
                 for i in range(1, len(self._rev2hsh)):
                     self._writerev(i, f)
         else:  # append incrementally
-            with open(self.path, b'ab') as f:
+            with open(self.path, 'ab') as f:
                 for i in range(self._lastmaxrev + 1, len(self._rev2hsh)):
                     self._writerev(i, f)
         self._lastmaxrev = self.maxrev
 
     def _load(self):
         """load state from file"""
         if not self.path:
             return
         # use local variables in a loop. CPython uses LOAD_FAST for them,
         # which is faster than both LOAD_CONST and LOAD_GLOBAL.
         flaglen = 1
         hshlen = _hshlen
-        with open(self.path, b'rb') as f:
+        with open(self.path, 'rb') as f:
             if f.read(len(self.HEADER)) != self.HEADER:
                 raise error.CorruptedFileError()
             self.clear(flush=False)
             while True:
                 buf = f.read(flaglen)
                 if not buf:
                     break
                 flag = ord(buf)
                 rev = len(self._rev2hsh)
                 if flag & renameflag:
                     path = self._readcstr(f)
                     self._renamerevs.append(rev)
                     self._renamepaths.append(path)
                 hsh = f.read(hshlen)
                 if len(hsh) != hshlen:
                     raise error.CorruptedFileError()
                 self._hsh2rev[hsh] = rev
                 self._rev2flag.append(flag)
                 self._rev2hsh.append(hsh)
         self._lastmaxrev = self.maxrev
 
     def _writerev(self, rev, f):
         """append a revision data to file"""
         flag = self._rev2flag[rev]
         hsh = self._rev2hsh[rev]
         f.write(struct.pack(b'B', flag))
         if flag & renameflag:
             path = self.rev2path(rev)
             if path is None:
                 raise error.CorruptedFileError(b'cannot find path for %s' % rev)
             f.write(path + b'\0')
         f.write(hsh)
 
     @staticmethod
     def _readcstr(f):
         """read a C-language-like '\0'-terminated string"""
         buf = b''
         while True:
             ch = f.read(1)
             if not ch:  # unexpected eof
                 raise error.CorruptedFileError()
             if ch == b'\0':
                 break
             buf += ch
         return buf
 
     def __contains__(self, f):
         """(fctx or (node, path)) -> bool.
         test if (node, path) is in the map, and is not in a side branch.
         f can be either a tuple of (node, path), or a fctx.
         """
         if isinstance(f, tuple):  # f: (node, path)
             hsh, path = f
         else:  # f: fctx
             hsh, path = f.node(), f.path()
         rev = self.hsh2rev(hsh)
         if rev is None:
             return False
         if path is not None and path != self.rev2path(rev):
             return False
         return (self.rev2flag(rev) & sidebranchflag) == 0
 
 
 def getlastnode(path):
     """return the last hash in a revmap, without loading its full content.
     this is equivalent to `m = revmap(path); m.rev2hsh(m.maxrev)`, but faster.
     """
     hsh = None
     try:
-        with open(path, b'rb') as f:
+        with open(path, 'rb') as f:
             f.seek(-_hshlen, io.SEEK_END)
             if f.tell() > len(revmap.HEADER):
                 hsh = f.read(_hshlen)
     except IOError:
         pass
     return hsh